From e6b770a8f002992ae59b5714bb91f063ef0921a0 Mon Sep 17 00:00:00 2001 From: Alberto Perdomo Date: Fri, 12 Jun 2026 19:46:19 +0100 Subject: [PATCH 1/6] fix: Make MultiModalFeatures per-prompt in TokenizedPrompt Signed-off-by: Alberto Perdomo --- .../interface/requesthandling/types.go | 14 +++++------ .../dataproducer/multimodal/producer.go | 10 ++++---- .../dataproducer/multimodal/producer_test.go | 6 ++--- .../preciseprefixcache/blockkeys.go | 8 +++---- .../preciseprefixcache/producer_test.go | 8 +++---- .../dataproducer/tokenizer/estimate.go | 12 ++++++++-- .../dataproducer/tokenizer/estimate_test.go | 20 +++++++++------- .../dataproducer/tokenizer/tokenizer.go | 9 +++---- .../dataproducer/tokenizer/tokenizer_test.go | 24 ++++++++++--------- .../parsers/vllmgrpc/vllmgrpc.go | 4 ++-- .../parsers/vllmgrpc/vllmgrpc_test.go | 8 +++---- .../disagg/disagg_profile_handler_test.go | 4 ++-- .../disagg/multimodal_helpers.go | 12 +++++++--- .../preciseprefixcache/legacy_producer.go | 4 ++-- 14 files changed, 82 insertions(+), 61 deletions(-) diff --git a/pkg/epp/framework/interface/requesthandling/types.go b/pkg/epp/framework/interface/requesthandling/types.go index c6a45cbbf8..4882a5f046 100644 --- a/pkg/epp/framework/interface/requesthandling/types.go +++ b/pkg/epp/framework/interface/requesthandling/types.go @@ -178,11 +178,10 @@ type TokenizedPrompt struct { // 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 + // MultiModalFeatures holds multimodal items per prompt, indexed in + // lockstep with PerPromptTokens. Single-prompt requests use a length-1 + // outer slice. Nil if the prompt contains no multimodal content. + MultiModalFeatures [][]MultiModalFeature // CacheSalt isolates prefix caches across requests. Populated by the token-producer. CacheSalt string } @@ -207,9 +206,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 corresponding PerPromptTokens entry. 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 } diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/multimodal/producer.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/multimodal/producer.go index b11dea3302..920906fd8d 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/multimodal/producer.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/multimodal/producer.go @@ -302,11 +302,13 @@ func ExtractMMItems(request *scheduling.InferenceRequest) []attrmm.MatchItem { } itemsByHash := map[string]attrmm.MatchItem{} - for _, feature := range request.Body.TokenizedPrompt.MultiModalFeatures { - if feature.Hash == "" { - continue + for _, perPrompt := range request.Body.TokenizedPrompt.MultiModalFeatures { + for _, feature := range perPrompt { + if feature.Hash == "" { + continue + } + addItem(itemsByHash, feature.Hash, string(feature.Modality)) } - addItem(itemsByHash, feature.Hash, string(feature.Modality)) } return itemSlice(itemsByHash) } diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/multimodal/producer_test.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/multimodal/producer_test.go index 2d9e808c59..5b9c04984b 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/multimodal/producer_test.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/multimodal/producer_test.go @@ -59,11 +59,11 @@ func TestExtractMMItemsFromTokenizedPrompt(t *testing.T) { items := ExtractMMItems(&scheduling.InferenceRequest{ Body: &fwkrh.InferenceRequestBody{ TokenizedPrompt: &fwkrh.TokenizedPrompt{ - MultiModalFeatures: []fwkrh.MultiModalFeature{ + MultiModalFeatures: [][]fwkrh.MultiModalFeature{{ {Modality: fwkrh.ModalityImage, Hash: "image-a", Length: 576}, {Modality: fwkrh.ModalityImage, Hash: "image-b", Length: 0}, {Modality: fwkrh.ModalityImage, Hash: "image-a", Length: 144}, - }, + }}, }, }, }) @@ -268,7 +268,7 @@ func requestWithHashes(requestID string, hashToWeight map[string]int) *schedulin return &scheduling.InferenceRequest{ RequestID: requestID, Body: &fwkrh.InferenceRequestBody{ - TokenizedPrompt: &fwkrh.TokenizedPrompt{MultiModalFeatures: features}, + TokenizedPrompt: &fwkrh.TokenizedPrompt{MultiModalFeatures: [][]fwkrh.MultiModalFeature{features}}, }, } } diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/blockkeys.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/blockkeys.go index 6300212141..52530bdbbc 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/blockkeys.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/blockkeys.go @@ -51,15 +51,13 @@ func computeBlockKeys(ctx context.Context, idx kvCacheIndexer, } var result [][]kvblock.BlockHash - for _, tokens := range tp.PerPromptTokens { + for i, tokens := range tp.PerPromptTokens { if len(tokens) == 0 { continue } - // MM features apply only to single-prompt requests (chat); multi-prompt - // completions never carry multimodal content. var mmf []fwkrh.MultiModalFeature - if len(tp.PerPromptTokens) == 1 { - mmf = tp.MultiModalFeatures + if i < len(tp.MultiModalFeatures) { + mmf = tp.MultiModalFeatures[i] } keys, err := computeBlockKeysForTokens(ctx, idx, tokens, mmf, tp.CacheSalt, request.TargetModel, blockSizeTokens) if err != nil { diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/producer_test.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/producer_test.go index 757ca27d04..c04e201648 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/producer_test.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/producer_test.go @@ -472,9 +472,9 @@ func TestProduce_PassesMMExtraFeatures(t *testing.T) { Body: &fwkrh.InferenceRequestBody{ TokenizedPrompt: &fwkrh.TokenizedPrompt{ PerPromptTokens: [][]uint32{tokens}, - MultiModalFeatures: []fwkrh.MultiModalFeature{ + MultiModalFeatures: [][]fwkrh.MultiModalFeature{{ {Modality: fwkrh.ModalityImage, Hash: "abc", Offset: 2, Length: 4}, - }, + }}, }, }, } @@ -496,7 +496,7 @@ func TestProduce_FoldsCacheSalt(t *testing.T) { tests := []struct { name string - mm []fwkrh.MultiModalFeature + mm [][]fwkrh.MultiModalFeature want []kvblock.MMHash }{ { @@ -505,7 +505,7 @@ func TestProduce_FoldsCacheSalt(t *testing.T) { }, { name: "salt appended after mm hash", - mm: []fwkrh.MultiModalFeature{{Modality: fwkrh.ModalityImage, Hash: "abc", Offset: 2, Length: 4}}, + mm: [][]fwkrh.MultiModalFeature{{{Modality: fwkrh.ModalityImage, Hash: "abc", Offset: 2, Length: 4}}}, want: []kvblock.MMHash{{Hash: "abc"}, {Hash: "s3cr3t"}}, }, } diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/estimate.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/estimate.go index dc95bbe225..1ead228a24 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/estimate.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/estimate.go @@ -128,7 +128,11 @@ func (b estimateBackend) produce(ctx context.Context, body *fwkrh.InferenceReque // and report them as features. if body.ChatCompletions != nil { raw, features := b.chatCompletionsBytes(body.ChatCompletions, mmMetadataFromContext(ctx)) - return &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{packBytes(raw)}, MultiModalFeatures: features}, nil + tp := &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{packBytes(raw)}} + if features != nil { + tp.MultiModalFeatures = [][]fwkrh.MultiModalFeature{features} + } + return tp, nil } if body.Messages != nil { raw, features := b.messagesBytes(body.Messages) @@ -140,7 +144,11 @@ func (b estimateBackend) produce(ctx context.Context, body *fwkrh.InferenceReque "mmFeatureCount", len(features), "mmFeatures", features, ) - return &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{tokens}, MultiModalFeatures: features}, nil + tp := &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{tokens}} + if features != nil { + tp.MultiModalFeatures = [][]fwkrh.MultiModalFeature{features} + } + return tp, nil } if body.Completions != nil && len(body.Completions.Prompt.Strings) > 1 { diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/estimate_test.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/estimate_test.go index 35c583ede4..f78f943ab6 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/estimate_test.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/estimate_test.go @@ -122,7 +122,8 @@ func TestEstimateBackend_ChatImageFeature(t *testing.T) { tp, err := estimateBackend{}.produce(context.Background(), body) require.NoError(t, err) require.Len(t, tp.MultiModalFeatures, 1) - f := tp.MultiModalFeatures[0] + require.Len(t, tp.MultiModalFeatures[0], 1) + f := tp.MultiModalFeatures[0][0] assert.Equal(t, fwkrh.ModalityImage, f.Modality) assert.Equal(t, strconv.FormatUint(xxhash.Sum64String(pngBase64DataURL), 16), f.Hash) assert.Greater(t, f.Length, 1, "image length must be > 1 (placeholder weighting)") @@ -173,10 +174,10 @@ func TestEstimateBackend_ChatImageWeightingDistinct(t *testing.T) { // Non-decodable URL falls back to the default 640x360 resolution. def, err := estimateBackend{}.produce(context.Background(), chat("https://example.com/a.png")) require.NoError(t, err) - assert.Equal(t, (defaultImageWidth*defaultImageHeight)/imageTokenFactor, def.MultiModalFeatures[0].Length, "default image length") + assert.Equal(t, (defaultImageWidth*defaultImageHeight)/imageTokenFactor, def.MultiModalFeatures[0][0].Length, "default image length") small, err := estimateBackend{}.produce(context.Background(), chat(pngBase64DataURL)) require.NoError(t, err) - assert.NotEqual(t, def.MultiModalFeatures[0].Length, small.MultiModalFeatures[0].Length, "different images yielded identical placeholder counts") + assert.NotEqual(t, def.MultiModalFeatures[0][0].Length, small.MultiModalFeatures[0][0].Length, "different images yielded identical placeholder counts") } // chatImageBody builds a chat request carrying a single image_url block. @@ -195,7 +196,8 @@ func TestImageEstimator_StaticMode(t *testing.T) { tp, err := b.produce(context.Background(), chatImageBody(pngBase64DataURL)) require.NoError(t, err) require.Len(t, tp.MultiModalFeatures, 1) - assert.Equal(t, 7, tp.MultiModalFeatures[0].Length, "static image length") + require.Len(t, tp.MultiModalFeatures[0], 1) + assert.Equal(t, 7, tp.MultiModalFeatures[0][0].Length, "static image length") } // TestImageEstimator_CustomFactor asserts the dynamic factor knob changes the @@ -205,7 +207,7 @@ func TestImageEstimator_CustomFactor(t *testing.T) { // Non-decodable URL falls back to the default 640x360 resolution. tp, err := b.produce(context.Background(), chatImageBody("https://example.com/a.png")) require.NoError(t, err) - assert.Equal(t, (defaultImageWidth*defaultImageHeight)/2048, tp.MultiModalFeatures[0].Length, "custom-factor image length") + assert.Equal(t, (defaultImageWidth*defaultImageHeight)/2048, tp.MultiModalFeatures[0][0].Length, "custom-factor image length") } // TestImageEstimator_CustomDefaultResolution asserts the default-resolution knob @@ -216,7 +218,7 @@ func TestImageEstimator_CustomDefaultResolution(t *testing.T) { }})} tp, err := b.produce(context.Background(), chatImageBody("https://example.com/a.png")) require.NoError(t, err) - assert.Equal(t, (1024*1024)/imageTokenFactor, tp.MultiModalFeatures[0].Length, "custom default-resolution length") + assert.Equal(t, (1024*1024)/imageTokenFactor, tp.MultiModalFeatures[0][0].Length, "custom default-resolution length") } // chatVideoBody builds a chat request carrying a single video_url block. @@ -525,7 +527,8 @@ func TestEstimateBackend_MessagesImageFeature(t *testing.T) { tp, err := estimateBackend{}.produce(context.Background(), body) require.NoError(t, err) require.Len(t, tp.MultiModalFeatures, 1) - f := tp.MultiModalFeatures[0] + require.Len(t, tp.MultiModalFeatures[0], 1) + f := tp.MultiModalFeatures[0][0] assert.Equal(t, fwkrh.ModalityImage, f.Modality) assert.Equal(t, strconv.FormatUint(xxhash.Sum64String(pngBase64Raw), 16), f.Hash, "base64 source must hash by its raw payload") assert.Greater(t, f.Length, 1, "image length must be > 1 (placeholder weighting)") @@ -550,7 +553,8 @@ func TestEstimateBackend_MessagesURLImageKey(t *testing.T) { tp, err := estimateBackend{}.produce(context.Background(), body) require.NoError(t, err) require.Len(t, tp.MultiModalFeatures, 1) - assert.Equal(t, strconv.FormatUint(xxhash.Sum64String(url), 16), tp.MultiModalFeatures[0].Hash) + require.Len(t, tp.MultiModalFeatures[0], 1) + assert.Equal(t, strconv.FormatUint(xxhash.Sum64String(url), 16), tp.MultiModalFeatures[0][0].Hash) } // TestEstimateBackend_MessagesDeterministic asserts identical requests produce diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/tokenizer.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/tokenizer.go index 9ff7e61afc..b0b5460ea1 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/tokenizer.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/tokenizer.go @@ -462,9 +462,10 @@ func anthropicImageToURL(src *fwkrh.AnthropicImageSource) string { } // convertMMFeaturesToUpstream flattens the kv-cache map-shaped multimodal -// metadata into the upstream flat list, sorted by placeholder offset so -// consumers see items in prompt order. Returns nil when no content is present. -func convertMMFeaturesToUpstream(src *tokenization.MultiModalFeatures) []fwkrh.MultiModalFeature { +// metadata into a single-prompt per-prompt slice, sorted by placeholder +// offset so consumers see items in prompt order. Returns nil when no +// content is present. +func convertMMFeaturesToUpstream(src *tokenization.MultiModalFeatures) [][]fwkrh.MultiModalFeature { if src == nil || len(src.MMHashes) == 0 { return nil } @@ -492,7 +493,7 @@ func convertMMFeaturesToUpstream(src *tokenization.MultiModalFeatures) []fwkrh.M return nil } sort.Slice(items, func(i, j int) bool { return items[i].Offset < items[j].Offset }) - return items + return [][]fwkrh.MultiModalFeature{items} } // ConvertMMFeaturesFromUpstream regroups the flat list of multimodal features diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/tokenizer_test.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/tokenizer_test.go index 34c6035ee6..cc34b3669b 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/tokenizer_test.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/tokenizer_test.go @@ -167,13 +167,14 @@ func TestProduce_PopulatesTokenizedPrompt(t *testing.T) { require.NoError(t, p.Produce(context.Background(), req, nil)) require.NotNil(t, req.Body.TokenizedPrompt) assert.Equal(t, []uint32{1, 2, 3, 4}, req.Body.TokenizedPrompt.PerPromptTokens[0]) - require.Len(t, req.Body.TokenizedPrompt.MultiModalFeatures, 2) - - assert.Equal(t, 3, req.Body.TokenizedPrompt.MultiModalFeatures[0].Offset) - assert.Equal(t, "hash-a", req.Body.TokenizedPrompt.MultiModalFeatures[0].Hash) - assert.Equal(t, 20, req.Body.TokenizedPrompt.MultiModalFeatures[1].Offset) - assert.Equal(t, "hash-b", req.Body.TokenizedPrompt.MultiModalFeatures[1].Hash) - assert.Equal(t, fwkrh.ModalityImage, req.Body.TokenizedPrompt.MultiModalFeatures[0].Modality) + require.Len(t, req.Body.TokenizedPrompt.MultiModalFeatures, 1) + require.Len(t, req.Body.TokenizedPrompt.MultiModalFeatures[0], 2) + + assert.Equal(t, 3, req.Body.TokenizedPrompt.MultiModalFeatures[0][0].Offset) + assert.Equal(t, "hash-a", req.Body.TokenizedPrompt.MultiModalFeatures[0][0].Hash) + assert.Equal(t, 20, req.Body.TokenizedPrompt.MultiModalFeatures[0][1].Offset) + assert.Equal(t, "hash-b", req.Body.TokenizedPrompt.MultiModalFeatures[0][1].Hash) + assert.Equal(t, fwkrh.ModalityImage, req.Body.TokenizedPrompt.MultiModalFeatures[0][0].Modality) } func TestProduce_SkipsWhenAlreadyPopulated(t *testing.T) { @@ -367,10 +368,10 @@ func TestProduce_GenerateFlattensFeatures(t *testing.T) { require.NotNil(t, req.Body.TokenizedPrompt) assert.Equal(t, tokenIDs, req.Body.TokenizedPrompt.PerPromptTokens[0]) assert.Equal(t, - []fwkrh.MultiModalFeature{ + [][]fwkrh.MultiModalFeature{{ {Modality: fwkrh.ModalityImage, Hash: "abc123hash", Offset: 1, Length: 3}, {Modality: fwkrh.ModalityImage, Hash: "def456hash", Offset: 4, Length: 3}, - }, + }}, req.Body.TokenizedPrompt.MultiModalFeatures, ) } @@ -383,9 +384,10 @@ func TestConvertMMFeaturesRoundTrip(t *testing.T) { }, } upstream := convertMMFeaturesToUpstream(src) - require.Len(t, upstream, 2) + require.Len(t, upstream, 1) + require.Len(t, upstream[0], 2) - hashes, ranges := ConvertMMFeaturesFromUpstream(upstream) + hashes, ranges := ConvertMMFeaturesFromUpstream(upstream[0]) assert.Equal(t, []string{"h1", "h2"}, hashes["image"]) assert.Equal(t, []kvblock.PlaceholderRange{{Offset: 1, Length: 2}, {Offset: 10, Length: 3}}, diff --git a/pkg/epp/framework/plugins/requesthandling/parsers/vllmgrpc/vllmgrpc.go b/pkg/epp/framework/plugins/requesthandling/parsers/vllmgrpc/vllmgrpc.go index 353e6ce9e1..c53b11ac13 100644 --- a/pkg/epp/framework/plugins/requesthandling/parsers/vllmgrpc/vllmgrpc.go +++ b/pkg/epp/framework/plugins/requesthandling/parsers/vllmgrpc/vllmgrpc.go @@ -240,7 +240,7 @@ func convertToInferenceRequestBody(pbReq *pb.GenerateRequest) (*fwkrh.InferenceR return body, nil } -func convertMultiModalFeatures(mmInputs *pb.MultimodalInputs) []fwkrh.MultiModalFeature { +func convertMultiModalFeatures(mmInputs *pb.MultimodalInputs) [][]fwkrh.MultiModalFeature { if mmInputs == nil { return nil } @@ -270,7 +270,7 @@ func convertMultiModalFeatures(mmInputs *pb.MultimodalInputs) []fwkrh.MultiModal features = append(features, feature) } - return features + return [][]fwkrh.MultiModalFeature{features} } func convertEmbedToInferenceRequestBody(pbReq *pb.EmbedRequest) (*fwkrh.InferenceRequestBody, error) { diff --git a/pkg/epp/framework/plugins/requesthandling/parsers/vllmgrpc/vllmgrpc_test.go b/pkg/epp/framework/plugins/requesthandling/parsers/vllmgrpc/vllmgrpc_test.go index 5a1839baf5..bce12b2282 100644 --- a/pkg/epp/framework/plugins/requesthandling/parsers/vllmgrpc/vllmgrpc_test.go +++ b/pkg/epp/framework/plugins/requesthandling/parsers/vllmgrpc/vllmgrpc_test.go @@ -185,10 +185,10 @@ func TestVllmGRPCParser_ParseRequest(t *testing.T) { }, TokenizedPrompt: &fwkrh.TokenizedPrompt{ PerPromptTokens: [][]uint32{{101, 102, 103, 104, 105}}, - MultiModalFeatures: []fwkrh.MultiModalFeature{ + MultiModalFeatures: [][]fwkrh.MultiModalFeature{{ {Modality: fwkrh.ModalityImage, Hash: "hash-a", Offset: 1, Length: 2}, {Modality: fwkrh.ModalityImage, Hash: "hash-b", Offset: 4, Length: 1}, - }, + }}, }, }, }, @@ -233,10 +233,10 @@ func TestVllmGRPCParser_ParseRequest(t *testing.T) { }, TokenizedPrompt: &fwkrh.TokenizedPrompt{ PerPromptTokens: [][]uint32{{201, 202, 203, 204}}, - MultiModalFeatures: []fwkrh.MultiModalFeature{ + MultiModalFeatures: [][]fwkrh.MultiModalFeature{{ {Modality: fwkrh.ModalityImage, Hash: "hash-only", Offset: 0, Length: 1}, {Modality: fwkrh.ModalityImage, Hash: "", Offset: 2, Length: 2}, - }, + }}, }, }, }, diff --git a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/disagg_profile_handler_test.go b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/disagg_profile_handler_test.go index 16c182504d..54afb0007a 100644 --- a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/disagg_profile_handler_test.go +++ b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/disagg_profile_handler_test.go @@ -103,7 +103,7 @@ func chatRequest(hasImage, hasVideo, hasAudio bool) *scheduling.InferenceRequest }, } if len(features) > 0 { - body.TokenizedPrompt = &fwkrh.TokenizedPrompt{MultiModalFeatures: features} + body.TokenizedPrompt = &fwkrh.TokenizedPrompt{MultiModalFeatures: [][]fwkrh.MultiModalFeature{features}} } return &scheduling.InferenceRequest{Body: body} } @@ -173,7 +173,7 @@ func TestHasMultimodalContent(t *testing.T) { {"feature present", &scheduling.InferenceRequest{ Body: &fwkrh.InferenceRequestBody{ TokenizedPrompt: &fwkrh.TokenizedPrompt{ - MultiModalFeatures: []fwkrh.MultiModalFeature{{Modality: fwkrh.ModalityImage}}, + MultiModalFeatures: [][]fwkrh.MultiModalFeature{{{Modality: fwkrh.ModalityImage}}}, }, }, }, true}, diff --git a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/multimodal_helpers.go b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/multimodal_helpers.go index 874de7bce6..d6ff54642e 100644 --- a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/multimodal_helpers.go +++ b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/multimodal_helpers.go @@ -8,7 +8,13 @@ import ( // multimodal features. Detection is protocol-agnostic: it relies on the // token-producer plugin having populated TokenizedPrompt.MultiModalFeatures. func hasMultimodalContent(request *scheduling.InferenceRequest) bool { - return request != nil && request.Body != nil && - request.Body.TokenizedPrompt != nil && - len(request.Body.TokenizedPrompt.MultiModalFeatures) > 0 + if request == nil || request.Body == nil || request.Body.TokenizedPrompt == nil { + return false + } + for _, perPrompt := range request.Body.TokenizedPrompt.MultiModalFeatures { + if len(perPrompt) > 0 { + return true + } + } + return false } diff --git a/pkg/epp/framework/plugins/scheduling/scorer/preciseprefixcache/legacy_producer.go b/pkg/epp/framework/plugins/scheduling/scorer/preciseprefixcache/legacy_producer.go index bef8b68e4d..e439eef2c1 100644 --- a/pkg/epp/framework/plugins/scheduling/scorer/preciseprefixcache/legacy_producer.go +++ b/pkg/epp/framework/plugins/scheduling/scorer/preciseprefixcache/legacy_producer.go @@ -142,7 +142,7 @@ func (lp *legacyProducer) tokenizeRequest(request *scheduling.InferenceRequest) // flattenMMFeatures regroups the kvcache map-shaped multimodal metadata // into the upstream flat list expected on TokenizedPrompt, sorted by // placeholder offset so consumers see items in prompt order. -func flattenMMFeatures(src *tokenization.MultiModalFeatures) []fwkrh.MultiModalFeature { +func flattenMMFeatures(src *tokenization.MultiModalFeatures) [][]fwkrh.MultiModalFeature { if src == nil || len(src.MMHashes) == 0 { return nil } @@ -169,5 +169,5 @@ func flattenMMFeatures(src *tokenization.MultiModalFeatures) []fwkrh.MultiModalF return nil } sort.Slice(items, func(i, j int) bool { return items[i].Offset < items[j].Offset }) - return items + return [][]fwkrh.MultiModalFeature{items} } From 800f520c1cc5c77a21e7ac42fb3a044475f6d4a6 Mon Sep 17 00:00:00 2001 From: Alberto Perdomo Date: Mon, 15 Jun 2026 10:03:30 +0100 Subject: [PATCH 2/6] fix: Bundle per-prompt tokens and MM features into PromptTokens struct Signed-off-by: Alberto Perdomo --- .../flowcontrol/benchmark/benchmark_test.go | 2 +- pkg/epp/flowcontrol/integration_test.go | 12 +-- .../interface/requesthandling/types.go | 39 ++++---- .../framework/interface/scheduling/types.go | 2 +- .../concurrency/detector_test.go | 2 +- .../dataproducer/approximateprefix/plugin.go | 4 +- .../approximateprefix/plugin_test.go | 22 ++--- .../dataproducer/inflightload/producer.go | 4 +- .../inflightload/producer_test.go | 6 +- .../inflightload/token_estimator.go | 4 +- .../inflightload/token_estimator_test.go | 6 +- .../dataproducer/multimodal/producer.go | 10 +- .../dataproducer/multimodal/producer_test.go | 20 ++-- .../preciseprefixcache/blockkeys.go | 28 +++--- .../dataproducer/preciseprefixcache/doc.go | 2 +- .../preciseprefixcache/producer.go | 6 +- .../preciseprefixcache/producer_test.go | 43 ++++----- .../predictedlatency/dataproducer_hooks.go | 2 +- .../dataproducer/predictedlatency/plugin.go | 2 +- .../predictedlatency/plugin_test.go | 4 +- .../requestcontrol_hooks_test.go | 4 +- .../dataproducer/prefixhash/hashing.go | 6 +- .../dataproducer/prefixhash/hashing_test.go | 36 +++---- .../dataproducer/tokenizer/backend.go | 35 ++++--- .../dataproducer/tokenizer/estimate.go | 40 ++++---- .../dataproducer/tokenizer/estimate_test.go | 96 +++++++++---------- .../dataproducer/tokenizer/tokenizer.go | 25 +++-- .../dataproducer/tokenizer/tokenizer_test.go | 91 +++++++++--------- .../dataproducer/tokenizer/vllm_http_test.go | 8 +- .../parsers/vllmgrpc/vllmgrpc.go | 12 ++- .../parsers/vllmgrpc/vllmgrpc_test.go | 28 +++--- .../disagg/disagg_profile_handler.go | 2 +- .../disagg/disagg_profile_handler_test.go | 25 +++-- .../disagg/multimodal_helpers.go | 8 +- .../disagg/pd_profile_handler.go | 2 +- .../disagg/pd_profile_handler_test.go | 2 +- .../disagg/prefix_based_pd_decider.go | 2 +- .../disagg/prefix_based_pd_decider_test.go | 12 +-- .../profilehandler/disagg/scheduler_test.go | 4 +- .../context_length_aware.go | 12 +-- .../context_length_aware_test.go | 10 +- .../preciseprefixcache/legacy_producer.go | 22 +++-- .../precise_prefix_cache_test.go | 26 ++--- 43 files changed, 373 insertions(+), 355 deletions(-) diff --git a/pkg/epp/flowcontrol/benchmark/benchmark_test.go b/pkg/epp/flowcontrol/benchmark/benchmark_test.go index 6e832c1f41..81d3217b63 100644 --- a/pkg/epp/flowcontrol/benchmark/benchmark_test.go +++ b/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) diff --git a/pkg/epp/flowcontrol/integration_test.go b/pkg/epp/flowcontrol/integration_test.go index 9ab3bd1518..37d05fa813 100644 --- a/pkg/epp/flowcontrol/integration_test.go +++ b/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{ diff --git a/pkg/epp/framework/interface/requesthandling/types.go b/pkg/epp/framework/interface/requesthandling/types.go index 4882a5f046..fa36a32f7a 100644 --- a/pkg/epp/framework/interface/requesthandling/types.go +++ b/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,31 +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 multimodal items per prompt, indexed in - // lockstep with PerPromptTokens. Single-prompt requests use a length-1 - // outer slice. Nil if the prompt contains no multimodal content. - 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,7 +212,7 @@ type MultiModalFeature struct { // 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 the corresponding PerPromptTokens entry. + // in the owning PromptTokens.TokenIDs slice. Offset int // Length is the number of placeholder tokens this item occupies. Length int diff --git a/pkg/epp/framework/interface/scheduling/types.go b/pkg/epp/framework/interface/scheduling/types.go index 056c8e9029..a167601ec5 100644 --- a/pkg/epp/framework/interface/scheduling/types.go +++ b/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 diff --git a/pkg/epp/framework/plugins/flowcontrol/saturationdetector/concurrency/detector_test.go b/pkg/epp/framework/plugins/flowcontrol/saturationdetector/concurrency/detector_test.go index 4264d85443..95b1f9c66d 100644 --- a/pkg/epp/framework/plugins/flowcontrol/saturationdetector/concurrency/detector_test.go +++ b/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)}}}, }, } } diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/plugin.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/plugin.go index d8f0709d18..c2273d8634 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/plugin.go +++ b/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{}}, } } diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/plugin_test.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/plugin_test.go index 6d22e99079..50b858228e 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/plugin_test.go +++ b/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}}}, }, }, } diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/inflightload/producer.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/inflightload/producer.go index 9b5b27dce5..d4f3c14dcf 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/inflightload/producer.go +++ b/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{}, diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/inflightload/producer_test.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/inflightload/producer_test.go index ff16c93837..021d2218a4 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/inflightload/producer_test.go +++ b/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)}}, }, }, } diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/inflightload/token_estimator.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/inflightload/token_estimator.go index 8bafc5e415..91dcf81f8a 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/inflightload/token_estimator.go +++ b/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 diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/inflightload/token_estimator_test.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/inflightload/token_estimator_test.go index 51595591f7..55460105f1 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/inflightload/token_estimator_test.go +++ b/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, diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/multimodal/producer.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/multimodal/producer.go index 920906fd8d..bf74a3ec49 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/multimodal/producer.go +++ b/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,13 +297,13 @@ 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 _, perPrompt := range request.Body.TokenizedPrompt.MultiModalFeatures { - for _, feature := range perPrompt { + for _, p := range request.Body.TokenizedRequest.Prompts { + for _, feature := range p.MultiModalFeatures { if feature.Hash == "" { continue } diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/multimodal/producer_test.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/multimodal/producer_test.go index 5b9c04984b..fcc7d42095 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/multimodal/producer_test.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/multimodal/producer_test.go @@ -55,14 +55,16 @@ func TestFactory(t *testing.T) { require.Error(t, err) } -func TestExtractMMItemsFromTokenizedPrompt(t *testing.T) { +func TestExtractMMItemsFromTokenizedRequest(t *testing.T) { items := ExtractMMItems(&scheduling.InferenceRequest{ Body: &fwkrh.InferenceRequestBody{ - TokenizedPrompt: &fwkrh.TokenizedPrompt{ - MultiModalFeatures: [][]fwkrh.MultiModalFeature{{ - {Modality: fwkrh.ModalityImage, Hash: "image-a", Length: 576}, - {Modality: fwkrh.ModalityImage, Hash: "image-b", Length: 0}, - {Modality: fwkrh.ModalityImage, Hash: "image-a", Length: 144}, + TokenizedRequest: &fwkrh.TokenizedRequest{ + Prompts: []fwkrh.PromptTokens{{ + MultiModalFeatures: []fwkrh.MultiModalFeature{ + {Modality: fwkrh.ModalityImage, Hash: "image-a", Length: 576}, + {Modality: fwkrh.ModalityImage, Hash: "image-b", Length: 0}, + {Modality: fwkrh.ModalityImage, Hash: "image-a", Length: 144}, + }, }}, }, }, @@ -74,7 +76,7 @@ func TestExtractMMItemsFromTokenizedPrompt(t *testing.T) { }, items) } -func TestExtractMMItemsNilTokenizedPromptReturnsNil(t *testing.T) { +func TestExtractMMItemsNilTokenizedRequestReturnsNil(t *testing.T) { items := ExtractMMItems(&scheduling.InferenceRequest{ Body: &fwkrh.InferenceRequestBody{}, }) @@ -84,7 +86,7 @@ func TestExtractMMItemsNilTokenizedPromptReturnsNil(t *testing.T) { func TestExtractMMItemsEmptyMultiModalFeaturesReturnsNil(t *testing.T) { items := ExtractMMItems(&scheduling.InferenceRequest{ Body: &fwkrh.InferenceRequestBody{ - TokenizedPrompt: &fwkrh.TokenizedPrompt{}, + TokenizedRequest: &fwkrh.TokenizedRequest{}, }, }) assert.Nil(t, items) @@ -268,7 +270,7 @@ func requestWithHashes(requestID string, hashToWeight map[string]int) *schedulin return &scheduling.InferenceRequest{ RequestID: requestID, Body: &fwkrh.InferenceRequestBody{ - TokenizedPrompt: &fwkrh.TokenizedPrompt{MultiModalFeatures: [][]fwkrh.MultiModalFeature{features}}, + TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{MultiModalFeatures: features}}}, }, } } diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/blockkeys.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/blockkeys.go index 52530bdbbc..612e17552a 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/blockkeys.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/blockkeys.go @@ -32,34 +32,30 @@ type kvCacheIndexer interface { KVBlockIndex() kvblock.Index } -// computeBlockKeys hashes the request's TokenizedPrompt into KV-block keys. -// When PerPromptTokens has more than one entry (multi-prompt completions), -// each prompt is hashed independently so cross-prompt block adjacency (which -// never exists in the model server cache) is avoided. Single-prompt requests -// produce a length-1 outer slice. A non-empty CacheSalt is folded into each -// prompt's first block. Returns nil when the request carries no tokens or no -// prompt produces full KV blocks. +// computeBlockKeys hashes the request's TokenizedRequest into KV-block keys. +// When Prompts has more than one entry (multi-prompt completions), each prompt +// is hashed independently so cross-prompt block adjacency (which never exists +// in the model server cache) is avoided. Single-prompt requests produce a +// length-1 outer slice. A non-empty CacheSalt is folded into each prompt's +// first block. Returns nil when the request carries no tokens or no prompt +// produces full KV blocks. func computeBlockKeys(ctx context.Context, idx kvCacheIndexer, request *scheduling.InferenceRequest, blockSizeTokens int, ) ([][]kvblock.BlockHash, error) { if request == nil || request.Body == nil { return nil, nil } - tp := request.Body.TokenizedPrompt - if tp == nil || len(tp.PerPromptTokens) == 0 { + tp := request.Body.TokenizedRequest + if tp == nil || len(tp.Prompts) == 0 { return nil, nil } var result [][]kvblock.BlockHash - for i, tokens := range tp.PerPromptTokens { - if len(tokens) == 0 { + for _, p := range tp.Prompts { + if len(p.TokenIDs) == 0 { continue } - var mmf []fwkrh.MultiModalFeature - if i < len(tp.MultiModalFeatures) { - mmf = tp.MultiModalFeatures[i] - } - keys, err := computeBlockKeysForTokens(ctx, idx, tokens, mmf, tp.CacheSalt, request.TargetModel, blockSizeTokens) + keys, err := computeBlockKeysForTokens(ctx, idx, p.TokenIDs, p.MultiModalFeatures, tp.CacheSalt, request.TargetModel, blockSizeTokens) if err != nil { return nil, err } diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/doc.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/doc.go index f4adce9b57..6df6a09acc 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/doc.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/doc.go @@ -14,6 +14,6 @@ See the License for the specific language governing permissions and limitations under the License. */ -// Package preciseprefixcache hashes TokenizedPrompt into KV-block keys, +// Package preciseprefixcache hashes TokenizedRequest into KV-block keys, // looks them up in the index, and writes per-endpoint PrefixCacheMatchInfo. package preciseprefixcache diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/producer.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/producer.go index 75fe087ac3..206d66e803 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/producer.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/producer.go @@ -285,15 +285,15 @@ func (p *Producer) Produces() map[plugin.DataKey]any { return map[plugin.DataKey]any{p.dk: attrprefix.PrefixCacheMatchInfo{}} } -// Consumes declares the TokenizedPrompt dependency from token-producer so +// Consumes declares the TokenizedRequest dependency from token-producer so // the data-layer DAG orders tokenization before this producer runs. 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{}}, } } -// Produce hashes the request's TokenizedPrompt into KV-block keys, looks +// Produce hashes the request's TokenizedRequest into KV-block keys, looks // them up in the per-endpoint KV-block index, and writes PrefixCacheMatchInfo // to each candidate endpoint. No-op when the request carries no tokens. // With speculativeIndexing enabled, the computed block keys are stashed diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/producer_test.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/producer_test.go index c04e201648..984aa45c2a 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/producer_test.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/producer_test.go @@ -150,7 +150,7 @@ func newProducerWithIndexer(ctx context.Context, idx kvCacheIndexer, scorer kvca } // Tokens present → Produce hashes and writes per-endpoint match info. -func TestProduce_UsesTokenizedPrompt(t *testing.T) { +func TestProduce_UsesTokenizedRequest(t *testing.T) { ctx := utils.NewTestContext(t) tokens := []uint32{10, 20, 30, 40, 50} @@ -183,7 +183,7 @@ func TestProduce_UsesTokenizedPrompt(t *testing.T) { RequestID: "req-1", TargetModel: "test-model", Body: &fwkrh.InferenceRequestBody{ - TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{tokens}}, + TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: tokens}}}, }, } @@ -228,7 +228,7 @@ func TestProduce_NoTokens_NoOp(t *testing.T) { } // Empty TokenIDs → no-op. -func TestProduce_EmptyTokenizedPrompt_NoOp(t *testing.T) { +func TestProduce_EmptyTokenizedRequest_NoOp(t *testing.T) { ctx := utils.NewTestContext(t) idx := &fakeKVCacheIndexer{ computeFromTokens: func(_ context.Context, _ []uint32, _ string, _ []*kvblock.BlockExtraFeatures) ([]kvblock.BlockHash, error) { @@ -243,8 +243,8 @@ func TestProduce_EmptyTokenizedPrompt_NoOp(t *testing.T) { RequestID: "req-3", TargetModel: "test-model", Body: &fwkrh.InferenceRequestBody{ - Completions: &fwkrh.CompletionsRequest{Prompt: fwkrh.Prompt{Raw: "p"}}, - TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{{}}}, + Completions: &fwkrh.CompletionsRequest{Prompt: fwkrh.Prompt{Raw: "p"}}, + TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: []uint32{}}}}, }, } require.NoError(t, p.Produce(ctx, req, testEndpoints)) @@ -285,8 +285,8 @@ func TestProduce_MultiPromptEmptyBlockKeys_NoOp(t *testing.T) { RequestID: "req-multi-empty", TargetModel: "test-model", Body: &fwkrh.InferenceRequestBody{ - TokenizedPrompt: &fwkrh.TokenizedPrompt{ - PerPromptTokens: [][]uint32{promptA, promptB}, + TokenizedRequest: &fwkrh.TokenizedRequest{ + Prompts: []fwkrh.PromptTokens{{TokenIDs: promptA}, {TokenIDs: promptB}}, }, }, } @@ -343,8 +343,8 @@ func TestProduce_MultiPromptSkipsEmptyPromptKeys(t *testing.T) { RequestID: "req-multi-mixed", TargetModel: "test-model", Body: &fwkrh.InferenceRequestBody{ - TokenizedPrompt: &fwkrh.TokenizedPrompt{ - PerPromptTokens: [][]uint32{shortPrompt, fullPrompt}, + TokenizedRequest: &fwkrh.TokenizedRequest{ + Prompts: []fwkrh.PromptTokens{{TokenIDs: shortPrompt}, {TokenIDs: fullPrompt}}, }, }, } @@ -470,10 +470,10 @@ func TestProduce_PassesMMExtraFeatures(t *testing.T) { RequestID: "req-mm", TargetModel: "test-model", Body: &fwkrh.InferenceRequestBody{ - TokenizedPrompt: &fwkrh.TokenizedPrompt{ - PerPromptTokens: [][]uint32{tokens}, - MultiModalFeatures: [][]fwkrh.MultiModalFeature{{ - {Modality: fwkrh.ModalityImage, Hash: "abc", Offset: 2, Length: 4}, + TokenizedRequest: &fwkrh.TokenizedRequest{ + Prompts: []fwkrh.PromptTokens{{ + TokenIDs: tokens, + MultiModalFeatures: []fwkrh.MultiModalFeature{{Modality: fwkrh.ModalityImage, Hash: "abc", Offset: 2, Length: 4}}, }}, }, }, @@ -496,7 +496,7 @@ func TestProduce_FoldsCacheSalt(t *testing.T) { tests := []struct { name string - mm [][]fwkrh.MultiModalFeature + mm []fwkrh.MultiModalFeature want []kvblock.MMHash }{ { @@ -505,7 +505,7 @@ func TestProduce_FoldsCacheSalt(t *testing.T) { }, { name: "salt appended after mm hash", - mm: [][]fwkrh.MultiModalFeature{{{Modality: fwkrh.ModalityImage, Hash: "abc", Offset: 2, Length: 4}}}, + mm: []fwkrh.MultiModalFeature{{Modality: fwkrh.ModalityImage, Hash: "abc", Offset: 2, Length: 4}}, want: []kvblock.MMHash{{Hash: "abc"}, {Hash: "s3cr3t"}}, }, } @@ -526,10 +526,9 @@ func TestProduce_FoldsCacheSalt(t *testing.T) { RequestID: "req-salt", TargetModel: "test-model", Body: &fwkrh.InferenceRequestBody{ - TokenizedPrompt: &fwkrh.TokenizedPrompt{ - PerPromptTokens: [][]uint32{tokens}, - MultiModalFeatures: tc.mm, - CacheSalt: "s3cr3t", + TokenizedRequest: &fwkrh.TokenizedRequest{ + Prompts: []fwkrh.PromptTokens{{TokenIDs: tokens, MultiModalFeatures: tc.mm}}, + CacheSalt: "s3cr3t", }, }, } @@ -564,7 +563,7 @@ func TestProduce_NoCacheSalt_NoExtraFeatures(t *testing.T) { RequestID: "req-nosalt", TargetModel: "test-model", Body: &fwkrh.InferenceRequestBody{ - TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{tokens}}, + TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: tokens}}}, }, } @@ -610,7 +609,7 @@ func TestProduces_DeclaresPrefixCacheMatchInfo(t *testing.T) { require.True(t, ok) } -func TestConsumes_DeclaresTokenizedPrompt(t *testing.T) { +func TestConsumes_DeclaresTokenizedRequest(t *testing.T) { p := &Producer{typedName: plugin.TypedName{Type: PluginType, Name: "x"}} expected := plugin.NewDataKey("TokenizedPrompt", "token-producer") _, ok := p.Consumes().Required[expected] @@ -661,7 +660,7 @@ func TestNew_BlockSizeFlowsViaTokenProcessor(t *testing.T) { RequestID: "r", TargetModel: "m", Body: &fwkrh.InferenceRequestBody{ - TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{tokens}}, + TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: tokens}}}, }, } require.NoError(t, p.Produce(ctx, req, []scheduling.Endpoint{endpoint})) diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/dataproducer_hooks.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/dataproducer_hooks.go index f1c556de2f..d66a99f605 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/dataproducer_hooks.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/dataproducer_hooks.go @@ -125,7 +125,7 @@ func (pl *PredictedLatency) Consumes() plugin.DataDependencies { Required: map[plugin.DataKey]any{ pl.prefixMatchDataKey: attrprefix.PrefixCacheMatchInfo{}, pl.inFlightLoadDataKey: attrconcurrency.InFlightLoad{}, - tokenproducer.TokenizedPromptDataKey: fwksched.TokenizedPrompt{}, + tokenproducer.TokenizedPromptDataKey: fwksched.TokenizedRequest{}, }, } } diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/plugin.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/plugin.go index e957209bdf..039ca1647c 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/plugin.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/plugin.go @@ -401,7 +401,7 @@ type predictedLatencyCtx struct { func newPredictedLatencyContext(request *fwksched.InferenceRequest) *predictedLatencyCtx { inputTokenCount := 0 if request.Body != nil { - if tp := request.Body.TokenizedPrompt; tp != nil { + if tp := request.Body.TokenizedRequest; tp != nil { inputTokenCount = tp.TokenCount() } } diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/plugin_test.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/plugin_test.go index 993c093c78..5ef4f0d55e 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/plugin_test.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/plugin_test.go @@ -147,7 +147,7 @@ func createTestInferenceRequest(reqID string, ttftSLO, tpotSLO float64) *fwksche Completions: &fwkrh.CompletionsRequest{ Prompt: fwkrh.Prompt{Raw: "test prompt"}, }, - TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{make([]uint32, 2)}}, + TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: make([]uint32, 2)}}}, }) } @@ -159,7 +159,7 @@ func createTestChatCompletionsInferenceRequest(reqID string, ttftSLO, tpotSLO fl {Role: "user", Content: fwkrh.Content{Raw: "Tell me a joke."}}, }, }, - TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{make([]uint32, 8)}}, + TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: make([]uint32, 8)}}}, }) } diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/requestcontrol_hooks_test.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/requestcontrol_hooks_test.go index 22660e7cc7..64de90d96f 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/requestcontrol_hooks_test.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/requestcontrol_hooks_test.go @@ -109,8 +109,8 @@ func TestNewPredictedLatencyContext_ChatCompletionsPrompt(t *testing.T) { func TestNewPredictedLatencyContext_GenerateUsesTokenIDCount(t *testing.T) { request := createTestInferenceRequestWithBody("test-generate", 1.0, 0.05, &fwkrh.InferenceRequestBody{ - Generate: &fwkrh.GenerateRequest{TokenIDs: []uint32{1, 2, 3, 4, 5}}, - TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{make([]uint32, 5)}}, + Generate: &fwkrh.GenerateRequest{TokenIDs: []uint32{1, 2, 3, 4, 5}}, + TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: make([]uint32, 5)}}}, }) ctx := newPredictedLatencyContext(request) diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/prefixhash/hashing.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/prefixhash/hashing.go index d01438b01c..1fbac26adc 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/prefixhash/hashing.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/prefixhash/hashing.go @@ -64,9 +64,9 @@ func GetBlockHashes(ctx context.Context, request *scheduling.InferenceRequest, b return nil } - tp := request.Body.TokenizedPrompt + tp := request.Body.TokenizedRequest if tp == nil || tp.TokenCount() == 0 { - loggerDebug.Info("TokenizedPrompt is empty, skipping hashing") + loggerDebug.Info("TokenizedRequest is empty, skipping hashing") return nil } @@ -92,7 +92,7 @@ func computeBlockHashes(seq iter.Seq[HashBlock], request *scheduling.InferenceRe h := xxhash.New() // Different models should have different hashes even with the same body. _, _ = h.Write([]byte(request.TargetModel)) - if cacheSalt := request.Body.TokenizedPrompt.CacheSalt; cacheSalt != "" { + if cacheSalt := request.Body.TokenizedRequest.CacheSalt; cacheSalt != "" { _, _ = h.Write([]byte(cacheSalt)) } diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/prefixhash/hashing_test.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/prefixhash/hashing_test.go index 37d8a239f9..eabde7e5df 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/prefixhash/hashing_test.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/prefixhash/hashing_test.go @@ -87,11 +87,11 @@ func TestGetBlockHashes(t *testing.T) { expectedBlocks int }{ { - name: "TokenizedPrompt", + name: "TokenizedRequest", request: &fwksched.InferenceRequest{ Body: &fwkrh.InferenceRequestBody{ - TokenizedPrompt: &fwkrh.TokenizedPrompt{ - PerPromptTokens: [][]uint32{{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}}, + TokenizedRequest: &fwkrh.TokenizedRequest{ + Prompts: []fwkrh.PromptTokens{{TokenIDs: []uint32{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}}}, }, }, }, @@ -99,7 +99,7 @@ func TestGetBlockHashes(t *testing.T) { expectedBlocks: 3, }, { - name: "MissingTokenizedPrompt", + name: "MissingTokenizedRequest", request: &fwksched.InferenceRequest{ Body: &fwkrh.InferenceRequestBody{}, }, @@ -110,7 +110,7 @@ func TestGetBlockHashes(t *testing.T) { name: "EmptyTokenIDs", request: &fwksched.InferenceRequest{ Body: &fwkrh.InferenceRequestBody{ - TokenizedPrompt: &fwkrh.TokenizedPrompt{}, + TokenizedRequest: &fwkrh.TokenizedRequest{}, }, }, blockSizeTokens: 4, @@ -139,9 +139,9 @@ func TestGetBlockHashes(t *testing.T) { func TestGetBlockHashesCacheSalt(t *testing.T) { body := func(salt string) *fwkrh.InferenceRequestBody { return &fwkrh.InferenceRequestBody{ - TokenizedPrompt: &fwkrh.TokenizedPrompt{ - PerPromptTokens: [][]uint32{{1, 2, 3, 4}}, - CacheSalt: salt, + TokenizedRequest: &fwkrh.TokenizedRequest{ + Prompts: []fwkrh.PromptTokens{{TokenIDs: []uint32{1, 2, 3, 4}}}, + CacheSalt: salt, }, } } @@ -160,28 +160,28 @@ func TestGetBlockHashesCacheSalt(t *testing.T) { func TestGetBlockHashes_MultiPrompt(t *testing.T) { tests := []struct { name string - perPromptTokens [][]uint32 + prompts []fwkrh.PromptTokens blockSizeTokens int expectedPrompts int expectedBlocksPerPrompt []int }{ { name: "TwoPrompts", - perPromptTokens: [][]uint32{{1, 2, 3, 4}, {5, 6, 7, 8}}, + prompts: []fwkrh.PromptTokens{{TokenIDs: []uint32{1, 2, 3, 4}}, {TokenIDs: []uint32{5, 6, 7, 8}}}, blockSizeTokens: 2, expectedPrompts: 2, expectedBlocksPerPrompt: []int{2, 2}, }, { name: "ThreePromptsUnevenLengths", - perPromptTokens: [][]uint32{{1, 2, 3}, {4, 5}, {6}}, + prompts: []fwkrh.PromptTokens{{TokenIDs: []uint32{1, 2, 3}}, {TokenIDs: []uint32{4, 5}}, {TokenIDs: []uint32{6}}}, blockSizeTokens: 2, expectedPrompts: 3, expectedBlocksPerPrompt: []int{2, 1, 1}, }, { name: "EmptyPromptSkipped", - perPromptTokens: [][]uint32{{1, 2}, {}, {3, 4}}, + prompts: []fwkrh.PromptTokens{{TokenIDs: []uint32{1, 2}}, {TokenIDs: []uint32{}}, {TokenIDs: []uint32{3, 4}}}, blockSizeTokens: 2, expectedPrompts: 2, expectedBlocksPerPrompt: []int{1, 1}, @@ -193,8 +193,8 @@ func TestGetBlockHashes_MultiPrompt(t *testing.T) { request := &fwksched.InferenceRequest{ TargetModel: "model", Body: &fwkrh.InferenceRequestBody{ - TokenizedPrompt: &fwkrh.TokenizedPrompt{ - PerPromptTokens: tt.perPromptTokens, + TokenizedRequest: &fwkrh.TokenizedRequest{ + Prompts: tt.prompts, }, }, } @@ -211,16 +211,16 @@ func TestGetBlockHashes_MultiPromptHashIndependence(t *testing.T) { multiPrompt := &fwksched.InferenceRequest{ TargetModel: "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}}}, }, }, } singlePrompt := &fwksched.InferenceRequest{ TargetModel: "model", Body: &fwkrh.InferenceRequestBody{ - TokenizedPrompt: &fwkrh.TokenizedPrompt{ - PerPromptTokens: [][]uint32{{1, 2, 3, 4}}, + TokenizedRequest: &fwkrh.TokenizedRequest{ + Prompts: []fwkrh.PromptTokens{{TokenIDs: []uint32{1, 2, 3, 4}}}, }, }, } diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/backend.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/backend.go index 718e414f50..2cf30c254e 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/backend.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/backend.go @@ -29,10 +29,10 @@ import ( fwkrh "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/requesthandling" ) -// tokenInputProducer turns a request body into a TokenizedPrompt. Backends vary +// tokenInputProducer turns a request body into a TokenizedRequest. Backends vary // in fidelity (render vs estimate); callers never branch on which produced it. type tokenInputProducer interface { - produce(ctx context.Context, body *fwkrh.InferenceRequestBody) (*fwkrh.TokenizedPrompt, error) + produce(ctx context.Context, body *fwkrh.InferenceRequestBody) (*fwkrh.TokenizedRequest, error) } // timeoutAware is implemented by backends (and the tokenizers they wrap) whose @@ -104,11 +104,11 @@ type renderBackend struct { tk tokenizer } -func (b renderBackend) produce(ctx context.Context, body *fwkrh.InferenceRequestBody) (*fwkrh.TokenizedPrompt, error) { +func (b renderBackend) produce(ctx context.Context, body *fwkrh.InferenceRequestBody) (*fwkrh.TokenizedRequest, error) { switch { case body.Completions != nil: if ids := body.Completions.Prompt.TokenIDs; len(ids) > 0 { - return &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{ids}}, nil + return &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: ids}}}, nil } return b.renderCompletions(ctx, body) case body.ChatCompletions != nil: @@ -116,21 +116,24 @@ func (b renderBackend) produce(ctx context.Context, body *fwkrh.InferenceRequest if err != nil { return nil, fmt.Errorf("tokenization failed: %w", err) } - return &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{tokenIDs}, MultiModalFeatures: convertMMFeaturesToUpstream(mmFeatures)}, nil + return &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{ + TokenIDs: tokenIDs, + MultiModalFeatures: convertMMFeaturesToUpstream(mmFeatures), + }}}, nil case body.Messages != nil: tokenIDs, mmFeatures, err := b.tk.RenderChat(ctx, messagesPayload(body)) if err != nil { return nil, fmt.Errorf("tokenization failed: %w", err) } - return &fwkrh.TokenizedPrompt{ - PerPromptTokens: [][]uint32{tokenIDs}, + return &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{ + TokenIDs: tokenIDs, MultiModalFeatures: convertMMFeaturesToUpstream(mmFeatures), - }, nil + }}}, nil case body.Generate != nil: - return &fwkrh.TokenizedPrompt{ - PerPromptTokens: [][]uint32{body.Generate.TokenIDs}, + return &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{ + TokenIDs: body.Generate.TokenIDs, MultiModalFeatures: convertMMFeaturesToUpstream(body.Generate.Features), - }, nil + }}}, nil default: return nil, errors.New("unsupported request body type, skipping tokenization") } @@ -181,7 +184,7 @@ func messagesPayload(body *fwkrh.InferenceRequestBody) fwkrh.RequestPayload { } // CacheSaltFromBody returns the cache salt from whichever protocol is populated. -// The protocol switch lives here so producers populate TokenizedPrompt.CacheSalt +// The protocol switch lives here so producers populate TokenizedRequest.CacheSalt // from one place and consumers read only that field. func CacheSaltFromBody(body *fwkrh.InferenceRequestBody) string { switch { @@ -207,10 +210,14 @@ func CacheSaltFromBody(body *fwkrh.InferenceRequestBody) string { // renderCompletions tokenizes a completions prompt via a single Render call. // completionsPayload builds the appropriate payload shape (single string or // string array), and the renderer returns the tokenized result. -func (b renderBackend) renderCompletions(ctx context.Context, body *fwkrh.InferenceRequestBody) (*fwkrh.TokenizedPrompt, error) { +func (b renderBackend) renderCompletions(ctx context.Context, body *fwkrh.InferenceRequestBody) (*fwkrh.TokenizedRequest, error) { allTokenIDs, _, err := b.tk.Render(ctx, completionsPayload(body)) if err != nil { return nil, fmt.Errorf("tokenization failed: %w", err) } - return &fwkrh.TokenizedPrompt{PerPromptTokens: allTokenIDs}, nil + prompts := make([]fwkrh.PromptTokens, len(allTokenIDs)) + for i, ids := range allTokenIDs { + prompts[i] = fwkrh.PromptTokens{TokenIDs: ids} + } + return &fwkrh.TokenizedRequest{Prompts: prompts}, nil } diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/estimate.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/estimate.go index 1ead228a24..c6729ed039 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/estimate.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/estimate.go @@ -108,31 +108,30 @@ func parseResolution(s string) (width, height int) { return w, h } -func (b estimateBackend) produce(ctx context.Context, body *fwkrh.InferenceRequestBody) (*fwkrh.TokenizedPrompt, error) { +func (b estimateBackend) produce(ctx context.Context, body *fwkrh.InferenceRequestBody) (*fwkrh.TokenizedRequest, error) { // Pre-tokenized inputs are already real tokens; pass them through unchanged // rather than byte-estimating. Token-ID inputs are valid for generate, // /v1/completions, and /v1/embeddings. switch { case body.Generate != nil: - return &fwkrh.TokenizedPrompt{ - PerPromptTokens: [][]uint32{body.Generate.TokenIDs}, + return &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{ + TokenIDs: body.Generate.TokenIDs, MultiModalFeatures: convertMMFeaturesToUpstream(body.Generate.Features), - }, nil + }}}, nil case body.Completions != nil && len(body.Completions.Prompt.TokenIDs) > 0: - return &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{body.Completions.Prompt.TokenIDs}}, nil + return &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: body.Completions.Prompt.TokenIDs}}}, nil case body.Embeddings != nil && len(body.Embeddings.Input.TokenIDs) > 0: - return &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{body.Embeddings.Input.TokenIDs}}, nil + return &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: body.Embeddings.Input.TokenIDs}}}, nil } // Chat and Anthropic messages fold multimodal placeholders into the stream // and report them as features. if body.ChatCompletions != nil { raw, features := b.chatCompletionsBytes(body.ChatCompletions, mmMetadataFromContext(ctx)) - tp := &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{packBytes(raw)}} - if features != nil { - tp.MultiModalFeatures = [][]fwkrh.MultiModalFeature{features} - } - return tp, nil + return &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{ + TokenIDs: packBytes(raw), + MultiModalFeatures: features, + }}}, nil } if body.Messages != nil { raw, features := b.messagesBytes(body.Messages) @@ -144,11 +143,10 @@ func (b estimateBackend) produce(ctx context.Context, body *fwkrh.InferenceReque "mmFeatureCount", len(features), "mmFeatures", features, ) - tp := &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{tokens}} - if features != nil { - tp.MultiModalFeatures = [][]fwkrh.MultiModalFeature{features} - } - return tp, nil + return &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{ + TokenIDs: tokens, + MultiModalFeatures: features, + }}}, nil } if body.Completions != nil && len(body.Completions.Prompt.Strings) > 1 { @@ -159,16 +157,16 @@ func (b estimateBackend) produce(ctx context.Context, body *fwkrh.InferenceReque if err != nil { return nil, err } - return &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{packBytes(raw)}}, nil + return &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: packBytes(raw)}}}, nil } -func estimateMultiStringCompletions(req *fwkrh.CompletionsRequest) (*fwkrh.TokenizedPrompt, error) { - allTokenIDs := make([][]uint32, 0, len(req.Prompt.Strings)) +func estimateMultiStringCompletions(req *fwkrh.CompletionsRequest) (*fwkrh.TokenizedRequest, error) { + prompts := make([]fwkrh.PromptTokens, 0, len(req.Prompt.Strings)) for _, s := range req.Prompt.Strings { ids := packBytes([]byte(s)) - allTokenIDs = append(allTokenIDs, ids) + prompts = append(prompts, fwkrh.PromptTokens{TokenIDs: ids}) } - return &fwkrh.TokenizedPrompt{PerPromptTokens: allTokenIDs}, nil + return &fwkrh.TokenizedRequest{Prompts: prompts}, nil } // estimateBytes serializes the user input of a non-chat request body to a byte diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/estimate_test.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/estimate_test.go index f78f943ab6..5e975cddb6 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/estimate_test.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/estimate_test.go @@ -58,7 +58,7 @@ func TestEstimateBackend_GeneratePassthrough(t *testing.T) { Generate: &fwkrh.GenerateRequest{TokenIDs: in}, }) require.NoError(t, err) - assert.Equal(t, in, tp.PerPromptTokens[0]) + assert.Equal(t, in, tp.Prompts[0].TokenIDs) } // TestEstimateBackend_CompletionsTokenIDsPassthrough asserts token-ID completions @@ -69,7 +69,7 @@ func TestEstimateBackend_CompletionsTokenIDsPassthrough(t *testing.T) { Completions: &fwkrh.CompletionsRequest{Prompt: fwkrh.Prompt{TokenIDs: in}}, }) require.NoError(t, err) - assert.Equal(t, in, tp.PerPromptTokens[0], "token IDs must pass through, not be byte-estimated") + assert.Equal(t, in, tp.Prompts[0].TokenIDs, "token IDs must pass through, not be byte-estimated") } // TestEstimateBackend_EmbeddingsTokenIDsPassthrough asserts token-ID embeddings @@ -80,7 +80,7 @@ func TestEstimateBackend_EmbeddingsTokenIDsPassthrough(t *testing.T) { Embeddings: &fwkrh.EmbeddingsRequest{Input: fwkrh.EmbeddingsInput{TokenIDs: in}}, }) require.NoError(t, err) - assert.Equal(t, in, tp.PerPromptTokens[0]) + assert.Equal(t, in, tp.Prompts[0].TokenIDs) } // TestEstimateBackend_CompletionsDeterministic asserts the same prompt produces @@ -93,10 +93,10 @@ func TestEstimateBackend_CompletionsDeterministic(t *testing.T) { require.NoError(t, err) b, err := estimateBackend{}.produce(context.Background(), body("hello world")) require.NoError(t, err) - assert.Equal(t, hashTokens(a.PerPromptTokens[0]), hashTokens(b.PerPromptTokens[0]), "same prompt produced different tokens") + assert.Equal(t, hashTokens(a.Prompts[0].TokenIDs), hashTokens(b.Prompts[0].TokenIDs), "same prompt produced different tokens") c, err := estimateBackend{}.produce(context.Background(), body("hello there")) require.NoError(t, err) - assert.NotEqual(t, hashTokens(a.PerPromptTokens[0]), hashTokens(c.PerPromptTokens[0]), "distinct prompts produced identical tokens") + assert.NotEqual(t, hashTokens(a.Prompts[0].TokenIDs), hashTokens(c.Prompts[0].TokenIDs), "distinct prompts produced identical tokens") } // pngBase64Raw is a 64x32 RGBA PNG (bare base64 payload), yielding @@ -121,14 +121,14 @@ func TestEstimateBackend_ChatImageFeature(t *testing.T) { } tp, err := estimateBackend{}.produce(context.Background(), body) require.NoError(t, err) - require.Len(t, tp.MultiModalFeatures, 1) - require.Len(t, tp.MultiModalFeatures[0], 1) - f := tp.MultiModalFeatures[0][0] + require.Len(t, tp.Prompts, 1) + require.Len(t, tp.Prompts[0].MultiModalFeatures, 1) + f := tp.Prompts[0].MultiModalFeatures[0] assert.Equal(t, fwkrh.ModalityImage, f.Modality) assert.Equal(t, strconv.FormatUint(xxhash.Sum64String(pngBase64DataURL), 16), f.Hash) assert.Greater(t, f.Length, 1, "image length must be > 1 (placeholder weighting)") assert.GreaterOrEqual(t, f.Offset, 0) - tokens := tp.PerPromptTokens[0] + tokens := tp.Prompts[0].TokenIDs assert.LessOrEqual(t, f.Offset+f.Length, len(tokens), "feature span [%d,%d) outside token stream of len %d", f.Offset, f.Offset+f.Length, len(tokens)) // Placeholder tokens are the URL hash repeated; verify the span carries weight. for i := f.Offset; i < f.Offset+f.Length; i++ { @@ -154,8 +154,8 @@ func TestEstimateBackend_ChatModalityLabels(t *testing.T) { t.Run(tc.name, func(t *testing.T) { tp, err := estimateBackend{}.produce(context.Background(), chat(tc.block)) require.NoError(t, err) - require.Len(t, tp.MultiModalFeatures, 1) - require.Equal(t, tc.want, tp.MultiModalFeatures[0].Modality) + require.Len(t, tp.Prompts[0].MultiModalFeatures, 1) + require.Equal(t, tc.want, tp.Prompts[0].MultiModalFeatures[0].Modality) }) } } @@ -174,10 +174,10 @@ func TestEstimateBackend_ChatImageWeightingDistinct(t *testing.T) { // Non-decodable URL falls back to the default 640x360 resolution. def, err := estimateBackend{}.produce(context.Background(), chat("https://example.com/a.png")) require.NoError(t, err) - assert.Equal(t, (defaultImageWidth*defaultImageHeight)/imageTokenFactor, def.MultiModalFeatures[0][0].Length, "default image length") + assert.Equal(t, (defaultImageWidth*defaultImageHeight)/imageTokenFactor, def.Prompts[0].MultiModalFeatures[0].Length, "default image length") small, err := estimateBackend{}.produce(context.Background(), chat(pngBase64DataURL)) require.NoError(t, err) - assert.NotEqual(t, def.MultiModalFeatures[0][0].Length, small.MultiModalFeatures[0][0].Length, "different images yielded identical placeholder counts") + assert.NotEqual(t, def.Prompts[0].MultiModalFeatures[0].Length, small.Prompts[0].MultiModalFeatures[0].Length, "different images yielded identical placeholder counts") } // chatImageBody builds a chat request carrying a single image_url block. @@ -195,9 +195,9 @@ func TestImageEstimator_StaticMode(t *testing.T) { b := estimateBackend{img: newImageEstimator(&estimateConfig{Image: &imageEstimateConfig{Mode: imageModeStatic, Static: &staticImageConfig{StaticToken: 7}}})} tp, err := b.produce(context.Background(), chatImageBody(pngBase64DataURL)) require.NoError(t, err) - require.Len(t, tp.MultiModalFeatures, 1) - require.Len(t, tp.MultiModalFeatures[0], 1) - assert.Equal(t, 7, tp.MultiModalFeatures[0][0].Length, "static image length") + require.Len(t, tp.Prompts, 1) + require.Len(t, tp.Prompts[0].MultiModalFeatures, 1) + assert.Equal(t, 7, tp.Prompts[0].MultiModalFeatures[0].Length, "static image length") } // TestImageEstimator_CustomFactor asserts the dynamic factor knob changes the @@ -207,7 +207,7 @@ func TestImageEstimator_CustomFactor(t *testing.T) { // Non-decodable URL falls back to the default 640x360 resolution. tp, err := b.produce(context.Background(), chatImageBody("https://example.com/a.png")) require.NoError(t, err) - assert.Equal(t, (defaultImageWidth*defaultImageHeight)/2048, tp.MultiModalFeatures[0][0].Length, "custom-factor image length") + assert.Equal(t, (defaultImageWidth*defaultImageHeight)/2048, tp.Prompts[0].MultiModalFeatures[0].Length, "custom-factor image length") } // TestImageEstimator_CustomDefaultResolution asserts the default-resolution knob @@ -218,7 +218,7 @@ func TestImageEstimator_CustomDefaultResolution(t *testing.T) { }})} tp, err := b.produce(context.Background(), chatImageBody("https://example.com/a.png")) require.NoError(t, err) - assert.Equal(t, (1024*1024)/imageTokenFactor, tp.MultiModalFeatures[0][0].Length, "custom default-resolution length") + assert.Equal(t, (1024*1024)/imageTokenFactor, tp.Prompts[0].MultiModalFeatures[0].Length, "custom default-resolution length") } // chatVideoBody builds a chat request carrying a single video_url block. @@ -526,9 +526,9 @@ func TestEstimateBackend_MessagesImageFeature(t *testing.T) { } tp, err := estimateBackend{}.produce(context.Background(), body) require.NoError(t, err) - require.Len(t, tp.MultiModalFeatures, 1) - require.Len(t, tp.MultiModalFeatures[0], 1) - f := tp.MultiModalFeatures[0][0] + require.Len(t, tp.Prompts, 1) + require.Len(t, tp.Prompts[0].MultiModalFeatures, 1) + f := tp.Prompts[0].MultiModalFeatures[0] assert.Equal(t, fwkrh.ModalityImage, f.Modality) assert.Equal(t, strconv.FormatUint(xxhash.Sum64String(pngBase64Raw), 16), f.Hash, "base64 source must hash by its raw payload") assert.Greater(t, f.Length, 1, "image length must be > 1 (placeholder weighting)") @@ -552,9 +552,9 @@ func TestEstimateBackend_MessagesURLImageKey(t *testing.T) { } tp, err := estimateBackend{}.produce(context.Background(), body) require.NoError(t, err) - require.Len(t, tp.MultiModalFeatures, 1) - require.Len(t, tp.MultiModalFeatures[0], 1) - assert.Equal(t, strconv.FormatUint(xxhash.Sum64String(url), 16), tp.MultiModalFeatures[0][0].Hash) + require.Len(t, tp.Prompts, 1) + require.Len(t, tp.Prompts[0].MultiModalFeatures, 1) + assert.Equal(t, strconv.FormatUint(xxhash.Sum64String(url), 16), tp.Prompts[0].MultiModalFeatures[0].Hash) } // TestEstimateBackend_MessagesDeterministic asserts identical requests produce @@ -574,10 +574,10 @@ func TestEstimateBackend_MessagesDeterministic(t *testing.T) { require.NoError(t, err) b, err := estimateBackend{}.produce(context.Background(), build("you are helpful", "hello world")) require.NoError(t, err) - assert.Equal(t, hashTokens(a.PerPromptTokens[0]), hashTokens(b.PerPromptTokens[0]), "identical messages requests produced different tokens") + assert.Equal(t, hashTokens(a.Prompts[0].TokenIDs), hashTokens(b.Prompts[0].TokenIDs), "identical messages requests produced different tokens") c, err := estimateBackend{}.produce(context.Background(), build("you are concise", "hello world")) require.NoError(t, err) - assert.NotEqual(t, hashTokens(a.PerPromptTokens[0]), hashTokens(c.PerPromptTokens[0]), "different system prompts produced identical tokens") + assert.NotEqual(t, hashTokens(a.Prompts[0].TokenIDs), hashTokens(c.Prompts[0].TokenIDs), "different system prompts produced identical tokens") } // TestEstimateBackend_ChatToolsBeforeSystem asserts the tools list is emitted @@ -602,9 +602,9 @@ func TestEstimateBackend_ChatToolsBeforeSystem(t *testing.T) { require.NoError(t, err) b, err := estimateBackend{}.produce(context.Background(), chat("you are a concise assistant")) require.NoError(t, err) - require.NotEqual(t, hashTokens(a.PerPromptTokens[0]), hashTokens(b.PerPromptTokens[0]), "streams identical, system was not applied") + require.NotEqual(t, hashTokens(a.Prompts[0].TokenIDs), hashTokens(b.Prompts[0].TokenIDs), "streams identical, system was not applied") for i := 0; i < sharedTokens; i++ { - assert.Equal(t, a.PerPromptTokens[0][i], b.PerPromptTokens[0][i], "token %d differs: tools should seed the prefix before system", i) + assert.Equal(t, a.Prompts[0].TokenIDs[i], b.Prompts[0].TokenIDs[i], "token %d differs: tools should seed the prefix before system", i) } } @@ -633,9 +633,9 @@ func TestEstimateBackend_MessagesToolsBeforeSystem(t *testing.T) { require.NoError(t, err) b, err := estimateBackend{}.produce(context.Background(), build("you are a concise assistant")) require.NoError(t, err) - require.NotEqual(t, hashTokens(a.PerPromptTokens[0]), hashTokens(b.PerPromptTokens[0]), "streams identical, system was not applied") + require.NotEqual(t, hashTokens(a.Prompts[0].TokenIDs), hashTokens(b.Prompts[0].TokenIDs), "streams identical, system was not applied") for i := 0; i < sharedTokens; i++ { - assert.Equal(t, a.PerPromptTokens[0][i], b.PerPromptTokens[0][i], "token %d differs: tools should seed the prefix before system", i) + assert.Equal(t, a.Prompts[0].TokenIDs[i], b.Prompts[0].TokenIDs[i], "token %d differs: tools should seed the prefix before system", i) } } @@ -656,14 +656,14 @@ func TestEstimateBackend_ChatToolsAffectPrefix(t *testing.T) { }} withTools, err := estimateBackend{}.produce(context.Background(), chat(weather)) require.NoError(t, err) - assert.NotEqual(t, hashTokens(noTools.PerPromptTokens[0]), hashTokens(withTools.PerPromptTokens[0]), "tools list was ignored by the prefix estimator") + assert.NotEqual(t, hashTokens(noTools.Prompts[0].TokenIDs), hashTokens(withTools.Prompts[0].TokenIDs), "tools list was ignored by the prefix estimator") stock := []any{map[string]any{ "type": "function", "function": map[string]any{"name": "get_stock_price"}, }} otherTools, err := estimateBackend{}.produce(context.Background(), chat(stock)) require.NoError(t, err) - assert.NotEqual(t, hashTokens(withTools.PerPromptTokens[0]), hashTokens(otherTools.PerPromptTokens[0]), "different tools lists produced identical tokens") + assert.NotEqual(t, hashTokens(withTools.Prompts[0].TokenIDs), hashTokens(otherTools.Prompts[0].TokenIDs), "different tools lists produced identical tokens") } // TestEstimateBackend_MessagesToolsAffectPrefix is the /v1/messages analog of @@ -686,7 +686,7 @@ func TestEstimateBackend_MessagesToolsAffectPrefix(t *testing.T) { }} withTools, err := estimateBackend{}.produce(context.Background(), build(weather)) require.NoError(t, err) - assert.NotEqual(t, hashTokens(noTools.PerPromptTokens[0]), hashTokens(withTools.PerPromptTokens[0]), "tools list was ignored by the messages prefix estimator") + assert.NotEqual(t, hashTokens(noTools.Prompts[0].TokenIDs), hashTokens(withTools.Prompts[0].TokenIDs), "tools list was ignored by the messages prefix estimator") stock := []any{map[string]any{ "name": "get_stock_price", "description": "Get a stock price", @@ -694,7 +694,7 @@ func TestEstimateBackend_MessagesToolsAffectPrefix(t *testing.T) { }} otherTools, err := estimateBackend{}.produce(context.Background(), build(stock)) require.NoError(t, err) - assert.NotEqual(t, hashTokens(withTools.PerPromptTokens[0]), hashTokens(otherTools.PerPromptTokens[0]), "different tools lists produced identical tokens") + assert.NotEqual(t, hashTokens(withTools.Prompts[0].TokenIDs), hashTokens(otherTools.Prompts[0].TokenIDs), "different tools lists produced identical tokens") } // TestEstimateBackend_NonChatNoFeatures asserts non-chat protocols carry no @@ -704,13 +704,13 @@ func TestEstimateBackend_NonChatNoFeatures(t *testing.T) { Completions: &fwkrh.CompletionsRequest{Prompt: fwkrh.Prompt{Raw: "hello"}}, }) require.NoError(t, err) - assert.Nil(t, tp.MultiModalFeatures, "non-chat features should be nil") + assert.Nil(t, tp.Prompts[0].MultiModalFeatures, "non-chat features should be nil") } -// TestEstimateBackend_MultiStringCompletionsPopulatesPerPromptTokens asserts -// that a multi-string completions prompt estimates each string independently -// and populates PerPromptTokens. -func TestEstimateBackend_MultiStringCompletionsPopulatesPerPromptTokens(t *testing.T) { +// TestEstimateBackend_MultiStringCompletionsPopulatesPrompts asserts that a +// multi-string completions prompt estimates each string independently and +// populates Prompts. +func TestEstimateBackend_MultiStringCompletionsPopulatesPrompts(t *testing.T) { tp, err := estimateBackend{}.produce(context.Background(), &fwkrh.InferenceRequestBody{ Completions: &fwkrh.CompletionsRequest{ Prompt: fwkrh.Prompt{Strings: []string{"hello world", "foo bar"}}, @@ -719,18 +719,18 @@ func TestEstimateBackend_MultiStringCompletionsPopulatesPerPromptTokens(t *testi if err != nil { t.Fatalf("produce: %v", err) } - if len(tp.PerPromptTokens) != 2 { - t.Fatalf("PerPromptTokens: got %d prompts, want 2", len(tp.PerPromptTokens)) + if len(tp.Prompts) != 2 { + t.Fatalf("Prompts: got %d prompts, want 2", len(tp.Prompts)) } - if tp.TokenCount() != len(tp.PerPromptTokens[0])+len(tp.PerPromptTokens[1]) { + if tp.TokenCount() != len(tp.Prompts[0].TokenIDs)+len(tp.Prompts[1].TokenIDs) { t.Errorf("flat TokenIDs length %d != sum of per-prompt lengths %d+%d", - tp.TokenCount(), len(tp.PerPromptTokens[0]), len(tp.PerPromptTokens[1])) + tp.TokenCount(), len(tp.Prompts[0].TokenIDs), len(tp.Prompts[1].TokenIDs)) } } -// TestEstimateBackend_SingleStringCompletionsSetsPerPromptTokens asserts that a -// single-element string array uses a length-1 PerPromptTokens slice. -func TestEstimateBackend_SingleStringCompletionsSetsPerPromptTokens(t *testing.T) { +// TestEstimateBackend_SingleStringCompletionsSetsPrompts asserts that a +// single-element string array uses a length-1 Prompts slice. +func TestEstimateBackend_SingleStringCompletionsSetsPrompts(t *testing.T) { tp, err := estimateBackend{}.produce(context.Background(), &fwkrh.InferenceRequestBody{ Completions: &fwkrh.CompletionsRequest{ Prompt: fwkrh.Prompt{Strings: []string{"hello world"}}, @@ -739,7 +739,7 @@ func TestEstimateBackend_SingleStringCompletionsSetsPerPromptTokens(t *testing.T if err != nil { t.Fatalf("produce: %v", err) } - if len(tp.PerPromptTokens) != 1 { - t.Errorf("single-string prompt should set length-1 PerPromptTokens, got %d", len(tp.PerPromptTokens)) + if len(tp.Prompts) != 1 { + t.Errorf("single-string prompt should set length-1 Prompts, got %d", len(tp.Prompts)) } } diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/tokenizer.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/tokenizer.go index b0b5460ea1..dd626e358e 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/tokenizer.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/tokenizer.go @@ -15,7 +15,7 @@ limitations under the License. */ // Package tokenizer provides a DataProducer plugin that tokenizes the request -// prompt and publishes the result on InferenceRequestBody.TokenizedPrompt for +// prompt and publishes the result on InferenceRequestBody.TokenizedRequest for // downstream consumers (scorers, filters, other data producers). package tokenizer @@ -301,7 +301,7 @@ func NewPlugin(ctx context.Context, name string, config *tokenizerPluginConfig) } // Plugin tokenizes the prompt in the incoming request and writes the result to -// InferenceRequestBody.TokenizedPrompt for downstream DataProducer / scoring plugins. +// InferenceRequestBody.TokenizedRequest for downstream DataProducer / scoring plugins. type Plugin struct { typedName plugin.TypedName backend tokenInputProducer @@ -321,7 +321,7 @@ func (p *Plugin) TypedName() plugin.TypedName { // Produces returns the data keys this plugin produces. func (p *Plugin) Produces() map[plugin.DataKey]any { - return map[plugin.DataKey]any{p.dk: fwkrh.TokenizedPrompt{}} + return map[plugin.DataKey]any{p.dk: fwkrh.TokenizedRequest{}} } // ProduceTimeout surfaces the backend's render timeout when it manages one, so @@ -334,18 +334,18 @@ func (p *Plugin) ProduceTimeout() time.Duration { return 0 } -// Produce derives the request's TokenizedPrompt via the configured backend and +// Produce derives the request's TokenizedRequest via the configured backend and // stores it on the body. Skips when one is already present; errors propagate to // the Director, which logs and continues. func (p *Plugin) Produce(ctx context.Context, request *scheduling.InferenceRequest, _ []scheduling.Endpoint) error { if request.Body == nil { return errors.New("request body is nil") } - if request.Body.TokenizedPrompt != nil { + if request.Body.TokenizedRequest != nil { // A parser (e.g. vLLM gRPC) may pre-populate tokens without a salt; // ensure cache-salt isolation still applies on the skip path. - if request.Body.TokenizedPrompt.CacheSalt == "" { - request.Body.TokenizedPrompt.CacheSalt = CacheSaltFromBody(request.Body) + if request.Body.TokenizedRequest.CacheSalt == "" { + request.Body.TokenizedRequest.CacheSalt = CacheSaltFromBody(request.Body) } return nil } @@ -359,7 +359,7 @@ func (p *Plugin) Produce(ctx context.Context, request *scheduling.InferenceReque return nil } tp.CacheSalt = CacheSaltFromBody(request.Body) - request.Body.TokenizedPrompt = tp + request.Body.TokenizedRequest = tp return nil } @@ -462,10 +462,9 @@ func anthropicImageToURL(src *fwkrh.AnthropicImageSource) string { } // convertMMFeaturesToUpstream flattens the kv-cache map-shaped multimodal -// metadata into a single-prompt per-prompt slice, sorted by placeholder -// offset so consumers see items in prompt order. Returns nil when no -// content is present. -func convertMMFeaturesToUpstream(src *tokenization.MultiModalFeatures) [][]fwkrh.MultiModalFeature { +// metadata into a flat list sorted by placeholder offset so consumers see +// items in prompt order. Returns nil when no content is present. +func convertMMFeaturesToUpstream(src *tokenization.MultiModalFeatures) []fwkrh.MultiModalFeature { if src == nil || len(src.MMHashes) == 0 { return nil } @@ -493,7 +492,7 @@ func convertMMFeaturesToUpstream(src *tokenization.MultiModalFeatures) [][]fwkrh return nil } sort.Slice(items, func(i, j int) bool { return items[i].Offset < items[j].Offset }) - return [][]fwkrh.MultiModalFeature{items} + return items } // ConvertMMFeaturesFromUpstream regroups the flat list of multimodal features diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/tokenizer_test.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/tokenizer_test.go index cc34b3669b..364061278d 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/tokenizer_test.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/tokenizer_test.go @@ -142,7 +142,7 @@ func TestPluginFactory_Validation(t *testing.T) { } } -func TestProduce_PopulatesTokenizedPrompt(t *testing.T) { +func TestProduce_PopulatesTokenizedRequest(t *testing.T) { mm := &tokenization.MultiModalFeatures{ MMHashes: map[string][]string{"image": {"hash-a", "hash-b"}}, MMPlaceholders: map[string][]kvblock.PlaceholderRange{ @@ -165,26 +165,26 @@ func TestProduce_PopulatesTokenizedPrompt(t *testing.T) { }, } require.NoError(t, p.Produce(context.Background(), req, nil)) - require.NotNil(t, req.Body.TokenizedPrompt) - assert.Equal(t, []uint32{1, 2, 3, 4}, req.Body.TokenizedPrompt.PerPromptTokens[0]) - require.Len(t, req.Body.TokenizedPrompt.MultiModalFeatures, 1) - require.Len(t, req.Body.TokenizedPrompt.MultiModalFeatures[0], 2) - - assert.Equal(t, 3, req.Body.TokenizedPrompt.MultiModalFeatures[0][0].Offset) - assert.Equal(t, "hash-a", req.Body.TokenizedPrompt.MultiModalFeatures[0][0].Hash) - assert.Equal(t, 20, req.Body.TokenizedPrompt.MultiModalFeatures[0][1].Offset) - assert.Equal(t, "hash-b", req.Body.TokenizedPrompt.MultiModalFeatures[0][1].Hash) - assert.Equal(t, fwkrh.ModalityImage, req.Body.TokenizedPrompt.MultiModalFeatures[0][0].Modality) + require.NotNil(t, req.Body.TokenizedRequest) + assert.Equal(t, []uint32{1, 2, 3, 4}, req.Body.TokenizedRequest.Prompts[0].TokenIDs) + require.Len(t, req.Body.TokenizedRequest.Prompts, 1) + require.Len(t, req.Body.TokenizedRequest.Prompts[0].MultiModalFeatures, 2) + + assert.Equal(t, 3, req.Body.TokenizedRequest.Prompts[0].MultiModalFeatures[0].Offset) + assert.Equal(t, "hash-a", req.Body.TokenizedRequest.Prompts[0].MultiModalFeatures[0].Hash) + assert.Equal(t, 20, req.Body.TokenizedRequest.Prompts[0].MultiModalFeatures[1].Offset) + assert.Equal(t, "hash-b", req.Body.TokenizedRequest.Prompts[0].MultiModalFeatures[1].Hash) + assert.Equal(t, fwkrh.ModalityImage, req.Body.TokenizedRequest.Prompts[0].MultiModalFeatures[0].Modality) } func TestProduce_SkipsWhenAlreadyPopulated(t *testing.T) { - existing := &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{{42}}} + existing := &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: []uint32{42}}}} p := newTestPlugin(&mockTokenizer{}) req := &scheduling.InferenceRequest{ - Body: &fwkrh.InferenceRequestBody{TokenizedPrompt: existing}, + Body: &fwkrh.InferenceRequestBody{TokenizedRequest: existing}, } require.NoError(t, p.Produce(context.Background(), req, nil)) - assert.Same(t, existing, req.Body.TokenizedPrompt) + assert.Same(t, existing, req.Body.TokenizedRequest) } func TestProduce_SetsCacheSaltOnSkipPath(t *testing.T) { @@ -194,18 +194,18 @@ func TestProduce_SetsCacheSaltOnSkipPath(t *testing.T) { return nil, nil, nil }, } - existing := &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{{1, 2, 3}}} + existing := &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: []uint32{1, 2, 3}}}} p := newTestPlugin(tok) req := &scheduling.InferenceRequest{ Body: &fwkrh.InferenceRequestBody{ - ChatCompletions: &fwkrh.ChatCompletionsRequest{CacheSalt: "tenant-x"}, - TokenizedPrompt: existing, + ChatCompletions: &fwkrh.ChatCompletionsRequest{CacheSalt: "tenant-x"}, + TokenizedRequest: existing, }, } require.NoError(t, p.Produce(context.Background(), req, nil)) - assert.Same(t, existing, req.Body.TokenizedPrompt) - assert.Equal(t, "tenant-x", req.Body.TokenizedPrompt.CacheSalt) - assert.Equal(t, []uint32{1, 2, 3}, req.Body.TokenizedPrompt.PerPromptTokens[0]) + assert.Same(t, existing, req.Body.TokenizedRequest) + assert.Equal(t, "tenant-x", req.Body.TokenizedRequest.CacheSalt) + assert.Equal(t, []uint32{1, 2, 3}, req.Body.TokenizedRequest.Prompts[0].TokenIDs) } func TestRenderBackend_CompletionsTokenIDsPassthrough(t *testing.T) { @@ -219,7 +219,7 @@ func TestRenderBackend_CompletionsTokenIDsPassthrough(t *testing.T) { Completions: &fwkrh.CompletionsRequest{Prompt: fwkrh.Prompt{TokenIDs: []uint32{5, 6, 7}}}, }) require.NoError(t, err) - assert.Equal(t, []uint32{5, 6, 7}, tp.PerPromptTokens[0]) + assert.Equal(t, []uint32{5, 6, 7}, tp.Prompts[0].TokenIDs) } func TestRenderBackend_CompletionsArrayPassesArrayPayload(t *testing.T) { @@ -237,7 +237,7 @@ func TestRenderBackend_CompletionsArrayPassesArrayPayload(t *testing.T) { Completions: &fwkrh.CompletionsRequest{Prompt: fwkrh.Prompt{Strings: []string{"alpha", "beta"}}}, }) require.NoError(t, err) - assert.Equal(t, [][]uint32{{1, 2}, {3}}, tp.PerPromptTokens) + assert.Equal(t, []fwkrh.PromptTokens{{TokenIDs: []uint32{1, 2}}, {TokenIDs: []uint32{3}}}, tp.Prompts) } func TestRenderBackend_CompletionsSingleArrayUsesPlainText(t *testing.T) { @@ -255,7 +255,7 @@ func TestRenderBackend_CompletionsSingleArrayUsesPlainText(t *testing.T) { }) require.NoError(t, err) assert.Equal(t, "alpha beta", got) - assert.Equal(t, [][]uint32{{1}}, tp.PerPromptTokens) + assert.Equal(t, []fwkrh.PromptTokens{{TokenIDs: []uint32{1}}}, tp.Prompts) } func TestProduce_NilBody(t *testing.T) { @@ -284,7 +284,7 @@ func TestProduce_TokenizerError(t *testing.T) { err := p.Produce(context.Background(), req, nil) require.Error(t, err) assert.Contains(t, err.Error(), "tokenization failed") - assert.Nil(t, req.Body.TokenizedPrompt) + assert.Nil(t, req.Body.TokenizedRequest) } func TestProduce_UnsupportedBodyType(t *testing.T) { @@ -297,7 +297,7 @@ func TestProduce_UnsupportedBodyType(t *testing.T) { err := p.Produce(context.Background(), req, nil) require.Error(t, err) assert.Contains(t, err.Error(), "unsupported request body type") - assert.Nil(t, req.Body.TokenizedPrompt) + assert.Nil(t, req.Body.TokenizedRequest) } func TestProduce_GenerateUsesPreTokenizedIDs(t *testing.T) { @@ -324,13 +324,13 @@ func TestProduce_GenerateUsesPreTokenizedIDs(t *testing.T) { } require.NoError(t, p.Produce(context.Background(), req, nil)) - require.NotNil(t, req.Body.TokenizedPrompt) - assert.Equal(t, tokenIDs, req.Body.TokenizedPrompt.PerPromptTokens[0]) - assert.Nil(t, req.Body.TokenizedPrompt.MultiModalFeatures) + require.NotNil(t, req.Body.TokenizedRequest) + assert.Equal(t, tokenIDs, req.Body.TokenizedRequest.Prompts[0].TokenIDs) + assert.Nil(t, req.Body.TokenizedRequest.Prompts[0].MultiModalFeatures) } func TestProduce_GenerateFlattensFeatures(t *testing.T) { - // Generate requests with multimodal features must populate TokenizedPrompt.MultiModalFeatures + // Generate requests with multimodal features must populate PromptTokens.MultiModalFeatures // in offset-sorted prompt order, so downstream prefix-cache scoring picks up image hashes. tok := &mockTokenizer{ renderFunc: func(_ fwkrh.RequestPayload) ([][]uint32, [][]tokenizerTypes.Offset, error) { @@ -365,14 +365,14 @@ func TestProduce_GenerateFlattensFeatures(t *testing.T) { } require.NoError(t, p.Produce(context.Background(), req, nil)) - require.NotNil(t, req.Body.TokenizedPrompt) - assert.Equal(t, tokenIDs, req.Body.TokenizedPrompt.PerPromptTokens[0]) + require.NotNil(t, req.Body.TokenizedRequest) + assert.Equal(t, tokenIDs, req.Body.TokenizedRequest.Prompts[0].TokenIDs) assert.Equal(t, - [][]fwkrh.MultiModalFeature{{ + []fwkrh.MultiModalFeature{ {Modality: fwkrh.ModalityImage, Hash: "abc123hash", Offset: 1, Length: 3}, {Modality: fwkrh.ModalityImage, Hash: "def456hash", Offset: 4, Length: 3}, - }}, - req.Body.TokenizedPrompt.MultiModalFeatures, + }, + req.Body.TokenizedRequest.Prompts[0].MultiModalFeatures, ) } @@ -384,10 +384,9 @@ func TestConvertMMFeaturesRoundTrip(t *testing.T) { }, } upstream := convertMMFeaturesToUpstream(src) - require.Len(t, upstream, 1) - require.Len(t, upstream[0], 2) + require.Len(t, upstream, 2) - hashes, ranges := ConvertMMFeaturesFromUpstream(upstream[0]) + hashes, ranges := ConvertMMFeaturesFromUpstream(upstream) assert.Equal(t, []string{"h1", "h2"}, hashes["image"]) assert.Equal(t, []kvblock.PlaceholderRange{{Offset: 1, Length: 2}, {Offset: 10, Length: 3}}, @@ -463,10 +462,10 @@ func TestProduce_StringArrayPrompt(t *testing.T) { }, } require.NoError(t, p.Produce(context.Background(), req, nil)) - require.NotNil(t, req.Body.TokenizedPrompt) - require.Len(t, req.Body.TokenizedPrompt.PerPromptTokens, 2) - assert.Equal(t, []uint32{10, 20, 30}, req.Body.TokenizedPrompt.PerPromptTokens[0]) - assert.Equal(t, []uint32{40, 50}, req.Body.TokenizedPrompt.PerPromptTokens[1]) + require.NotNil(t, req.Body.TokenizedRequest) + require.Len(t, req.Body.TokenizedRequest.Prompts, 2) + assert.Equal(t, []uint32{10, 20, 30}, req.Body.TokenizedRequest.Prompts[0].TokenIDs) + assert.Equal(t, []uint32{40, 50}, req.Body.TokenizedRequest.Prompts[1].TokenIDs) } func TestProduce_StringArrayPromptRenderError(t *testing.T) { @@ -487,7 +486,7 @@ func TestProduce_StringArrayPromptRenderError(t *testing.T) { err := p.Produce(context.Background(), req, nil) require.Error(t, err) assert.Contains(t, err.Error(), "tokenization failed") - assert.Nil(t, req.Body.TokenizedPrompt) + assert.Nil(t, req.Body.TokenizedRequest) } func TestProduce_StringArrayPromptDoesNotPublishEmptyTokenResult(t *testing.T) { @@ -506,10 +505,10 @@ func TestProduce_StringArrayPromptDoesNotPublishEmptyTokenResult(t *testing.T) { }, } require.NoError(t, p.Produce(context.Background(), req, nil)) - assert.Nil(t, req.Body.TokenizedPrompt) + assert.Nil(t, req.Body.TokenizedRequest) } -func TestProduce_SinglePromptSetsPerPromptTokens(t *testing.T) { +func TestProduce_SinglePromptSetsPrompts(t *testing.T) { tok := &mockTokenizer{ renderFunc: func(_ fwkrh.RequestPayload) ([][]uint32, [][]tokenizerTypes.Offset, error) { return [][]uint32{{10, 20, 30}}, nil, nil @@ -525,8 +524,8 @@ func TestProduce_SinglePromptSetsPerPromptTokens(t *testing.T) { }, } require.NoError(t, p.Produce(context.Background(), req, nil)) - require.NotNil(t, req.Body.TokenizedPrompt) - assert.Equal(t, [][]uint32{{10, 20, 30}}, req.Body.TokenizedPrompt.PerPromptTokens) + require.NotNil(t, req.Body.TokenizedRequest) + assert.Equal(t, []fwkrh.PromptTokens{{TokenIDs: []uint32{10, 20, 30}}}, req.Body.TokenizedRequest.Prompts) } func TestChatCompletionsToRenderChatRequest_MultimodalContent(t *testing.T) { diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/vllm_http_test.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/vllm_http_test.go index 69e7f0f431..7ae32c59fb 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/vllm_http_test.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/vllm_http_test.go @@ -96,8 +96,8 @@ func TestProduce_CompletionsVLLMHTTPUsesRawPayload(t *testing.T) { p := newTestPlugin(newHTTPRenderer(t, srv)) require.NoError(t, p.Produce(context.Background(), req, nil)) - require.NotNil(t, req.Body.TokenizedPrompt) - assert.Equal(t, []uint32{4, 5}, req.Body.TokenizedPrompt.PerPromptTokens[0]) + require.NotNil(t, req.Body.TokenizedRequest) + assert.Equal(t, []uint32{4, 5}, req.Body.TokenizedRequest.Prompts[0].TokenIDs) var sent map[string]any require.NoError(t, json.Unmarshal(cap.completions, &sent)) @@ -229,8 +229,8 @@ func TestProduce_ChatCompletionsVLLMHTTPUsesRawPayload(t *testing.T) { p := newTestPlugin(newHTTPRenderer(t, srv)) require.NoError(t, p.Produce(context.Background(), req, nil)) - require.NotNil(t, req.Body.TokenizedPrompt) - assert.Equal(t, []uint32{9, 10}, req.Body.TokenizedPrompt.PerPromptTokens[0]) + require.NotNil(t, req.Body.TokenizedRequest) + assert.Equal(t, []uint32{9, 10}, req.Body.TokenizedRequest.Prompts[0].TokenIDs) var sent map[string]any require.NoError(t, json.Unmarshal(cap.chat, &sent)) diff --git a/pkg/epp/framework/plugins/requesthandling/parsers/vllmgrpc/vllmgrpc.go b/pkg/epp/framework/plugins/requesthandling/parsers/vllmgrpc/vllmgrpc.go index c53b11ac13..7852ef57c9 100644 --- a/pkg/epp/framework/plugins/requesthandling/parsers/vllmgrpc/vllmgrpc.go +++ b/pkg/epp/framework/plugins/requesthandling/parsers/vllmgrpc/vllmgrpc.go @@ -223,9 +223,11 @@ func convertToInferenceRequestBody(pbReq *pb.GenerateRequest) (*fwkrh.InferenceR Prompt: fwkrh.Prompt{TokenIDs: copiedTokenIDsInt}, }, Payload: fwkrh.PayloadProto{Message: pbReq}, - TokenizedPrompt: &fwkrh.TokenizedPrompt{ - PerPromptTokens: [][]uint32{copiedTokenIDsInt}, - MultiModalFeatures: convertMultiModalFeatures(pbReq.GetMmInputs()), + TokenizedRequest: &fwkrh.TokenizedRequest{ + Prompts: []fwkrh.PromptTokens{{ + TokenIDs: copiedTokenIDsInt, + MultiModalFeatures: convertMultiModalFeatures(pbReq.GetMmInputs()), + }}, }, } default: @@ -240,7 +242,7 @@ func convertToInferenceRequestBody(pbReq *pb.GenerateRequest) (*fwkrh.InferenceR return body, nil } -func convertMultiModalFeatures(mmInputs *pb.MultimodalInputs) [][]fwkrh.MultiModalFeature { +func convertMultiModalFeatures(mmInputs *pb.MultimodalInputs) []fwkrh.MultiModalFeature { if mmInputs == nil { return nil } @@ -270,7 +272,7 @@ func convertMultiModalFeatures(mmInputs *pb.MultimodalInputs) [][]fwkrh.MultiMod features = append(features, feature) } - return [][]fwkrh.MultiModalFeature{features} + return features } func convertEmbedToInferenceRequestBody(pbReq *pb.EmbedRequest) (*fwkrh.InferenceRequestBody, error) { diff --git a/pkg/epp/framework/plugins/requesthandling/parsers/vllmgrpc/vllmgrpc_test.go b/pkg/epp/framework/plugins/requesthandling/parsers/vllmgrpc/vllmgrpc_test.go index bce12b2282..1aaeb283d8 100644 --- a/pkg/epp/framework/plugins/requesthandling/parsers/vllmgrpc/vllmgrpc_test.go +++ b/pkg/epp/framework/plugins/requesthandling/parsers/vllmgrpc/vllmgrpc_test.go @@ -139,8 +139,8 @@ func TestVllmGRPCParser_ParseRequest(t *testing.T) { }, }, }}, - TokenizedPrompt: &fwkrh.TokenizedPrompt{ - PerPromptTokens: [][]uint32{{11, 12, 13}}, + TokenizedRequest: &fwkrh.TokenizedRequest{ + Prompts: []fwkrh.PromptTokens{{TokenIDs: []uint32{11, 12, 13}}}, }, }, }, @@ -183,11 +183,13 @@ func TestVllmGRPCParser_ParseRequest(t *testing.T) { }, }, }, - TokenizedPrompt: &fwkrh.TokenizedPrompt{ - PerPromptTokens: [][]uint32{{101, 102, 103, 104, 105}}, - MultiModalFeatures: [][]fwkrh.MultiModalFeature{{ - {Modality: fwkrh.ModalityImage, Hash: "hash-a", Offset: 1, Length: 2}, - {Modality: fwkrh.ModalityImage, Hash: "hash-b", Offset: 4, Length: 1}, + TokenizedRequest: &fwkrh.TokenizedRequest{ + Prompts: []fwkrh.PromptTokens{{ + TokenIDs: []uint32{101, 102, 103, 104, 105}, + MultiModalFeatures: []fwkrh.MultiModalFeature{ + {Modality: fwkrh.ModalityImage, Hash: "hash-a", Offset: 1, Length: 2}, + {Modality: fwkrh.ModalityImage, Hash: "hash-b", Offset: 4, Length: 1}, + }, }}, }, }, @@ -231,11 +233,13 @@ func TestVllmGRPCParser_ParseRequest(t *testing.T) { }, }, }, - TokenizedPrompt: &fwkrh.TokenizedPrompt{ - PerPromptTokens: [][]uint32{{201, 202, 203, 204}}, - MultiModalFeatures: [][]fwkrh.MultiModalFeature{{ - {Modality: fwkrh.ModalityImage, Hash: "hash-only", Offset: 0, Length: 1}, - {Modality: fwkrh.ModalityImage, Hash: "", Offset: 2, Length: 2}, + TokenizedRequest: &fwkrh.TokenizedRequest{ + Prompts: []fwkrh.PromptTokens{{ + TokenIDs: []uint32{201, 202, 203, 204}, + MultiModalFeatures: []fwkrh.MultiModalFeature{ + {Modality: fwkrh.ModalityImage, Hash: "hash-only", Offset: 0, Length: 1}, + {Modality: fwkrh.ModalityImage, Hash: "", Offset: 2, Length: 2}, + }, }}, }, }, diff --git a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/disagg_profile_handler.go b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/disagg_profile_handler.go index 21a74b7573..8904055324 100644 --- a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/disagg_profile_handler.go +++ b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/disagg_profile_handler.go @@ -255,7 +255,7 @@ func (*Handler) Consumes() plugin.DataDependencies { return plugin.DataDependencies{ Required: map[plugin.DataKey]any{ attrprefix.PrefixCacheMatchInfoDataKey: attrprefix.PrefixCacheMatchInfo{}, - tokenproducer.TokenizedPromptDataKey: scheduling.TokenizedPrompt{}, + tokenproducer.TokenizedPromptDataKey: scheduling.TokenizedRequest{}, }, } } diff --git a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/disagg_profile_handler_test.go b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/disagg_profile_handler_test.go index 54afb0007a..98aab0877d 100644 --- a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/disagg_profile_handler_test.go +++ b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/disagg_profile_handler_test.go @@ -73,15 +73,15 @@ func profileNames(m map[string]scheduling.SchedulerProfile) []string { func completionsRequest(prompt string) *scheduling.InferenceRequest { return &scheduling.InferenceRequest{ Body: &fwkrh.InferenceRequestBody{ - Completions: &fwkrh.CompletionsRequest{Prompt: fwkrh.Prompt{Raw: prompt}}, - TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{make([]uint32, len(prompt)/averageCharactersPerToken)}}, + Completions: &fwkrh.CompletionsRequest{Prompt: fwkrh.Prompt{Raw: prompt}}, + TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: make([]uint32, len(prompt)/averageCharactersPerToken)}}}, }, } } // chatRequest builds a chat-completions InferenceRequest, populating the // tokenized prompt with one multimodal feature per requested modality so -// that multimodal detection (which reads TokenizedPrompt) is exercised. +// that multimodal detection (which reads TokenizedRequest) is exercised. func chatRequest(hasImage, hasVideo, hasAudio bool) *scheduling.InferenceRequest { blocks := []fwkrh.ContentBlock{{Type: "text", Text: "describe this"}} var features []fwkrh.MultiModalFeature @@ -103,7 +103,7 @@ func chatRequest(hasImage, hasVideo, hasAudio bool) *scheduling.InferenceRequest }, } if len(features) > 0 { - body.TokenizedPrompt = &fwkrh.TokenizedPrompt{MultiModalFeatures: [][]fwkrh.MultiModalFeature{features}} + body.TokenizedRequest = &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{MultiModalFeatures: features}}} } return &scheduling.InferenceRequest{Body: body} } @@ -113,10 +113,15 @@ func chatRequest(hasImage, hasVideo, hasAudio bool) *scheduling.InferenceRequest // any existing multimodal features. func withPrompt(req *scheduling.InferenceRequest, prompt string) *scheduling.InferenceRequest { req.Body.Completions = &fwkrh.CompletionsRequest{Prompt: fwkrh.Prompt{Raw: prompt}} - if req.Body.TokenizedPrompt == nil { - req.Body.TokenizedPrompt = &fwkrh.TokenizedPrompt{} + if req.Body.TokenizedRequest == nil { + req.Body.TokenizedRequest = &fwkrh.TokenizedRequest{} + } + tokenIDs := make([]uint32, len(prompt)/averageCharactersPerToken) + if len(req.Body.TokenizedRequest.Prompts) > 0 { + req.Body.TokenizedRequest.Prompts[0].TokenIDs = tokenIDs + } else { + req.Body.TokenizedRequest.Prompts = []fwkrh.PromptTokens{{TokenIDs: tokenIDs}} } - req.Body.TokenizedPrompt.PerPromptTokens = [][]uint32{make([]uint32, len(prompt)/averageCharactersPerToken)} return req } @@ -164,7 +169,7 @@ func TestHasMultimodalContent(t *testing.T) { {"nil body", &scheduling.InferenceRequest{Body: nil}, false}, {"nil tokenized prompt", &scheduling.InferenceRequest{Body: &fwkrh.InferenceRequestBody{}}, false}, {"empty multimodal features", &scheduling.InferenceRequest{ - Body: &fwkrh.InferenceRequestBody{TokenizedPrompt: &fwkrh.TokenizedPrompt{}}, + Body: &fwkrh.InferenceRequestBody{TokenizedRequest: &fwkrh.TokenizedRequest{}}, }, false}, {"text only", chatRequest(false, false, false), false}, {"image", chatRequest(true, false, false), true}, @@ -172,8 +177,8 @@ func TestHasMultimodalContent(t *testing.T) { {"audio", chatRequest(false, false, true), true}, {"feature present", &scheduling.InferenceRequest{ Body: &fwkrh.InferenceRequestBody{ - TokenizedPrompt: &fwkrh.TokenizedPrompt{ - MultiModalFeatures: [][]fwkrh.MultiModalFeature{{{Modality: fwkrh.ModalityImage}}}, + TokenizedRequest: &fwkrh.TokenizedRequest{ + Prompts: []fwkrh.PromptTokens{{MultiModalFeatures: []fwkrh.MultiModalFeature{{Modality: fwkrh.ModalityImage}}}}, }, }, }, true}, diff --git a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/multimodal_helpers.go b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/multimodal_helpers.go index d6ff54642e..bdeadd4c1c 100644 --- a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/multimodal_helpers.go +++ b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/multimodal_helpers.go @@ -6,13 +6,13 @@ import ( // hasMultimodalContent reports whether the tokenized prompt carries any // multimodal features. Detection is protocol-agnostic: it relies on the -// token-producer plugin having populated TokenizedPrompt.MultiModalFeatures. +// token-producer plugin having populated PromptTokens.MultiModalFeatures. func hasMultimodalContent(request *scheduling.InferenceRequest) bool { - if request == nil || request.Body == nil || request.Body.TokenizedPrompt == nil { + if request == nil || request.Body == nil || request.Body.TokenizedRequest == nil { return false } - for _, perPrompt := range request.Body.TokenizedPrompt.MultiModalFeatures { - if len(perPrompt) > 0 { + for _, p := range request.Body.TokenizedRequest.Prompts { + if len(p.MultiModalFeatures) > 0 { return true } } diff --git a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/pd_profile_handler.go b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/pd_profile_handler.go index 3065fbd3ad..d52a308a78 100644 --- a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/pd_profile_handler.go +++ b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/pd_profile_handler.go @@ -149,7 +149,7 @@ func (h *PdProfileHandler) Consumes() plugin.DataDependencies { return plugin.DataDependencies{ Required: map[plugin.DataKey]any{ h.dk: attrprefix.PrefixCacheMatchInfo{}, - tokenproducer.TokenizedPromptDataKey: scheduling.TokenizedPrompt{}, + tokenproducer.TokenizedPromptDataKey: scheduling.TokenizedRequest{}, }, } } diff --git a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/pd_profile_handler_test.go b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/pd_profile_handler_test.go index efde5d9860..cf4fefa0b1 100644 --- a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/pd_profile_handler_test.go +++ b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/pd_profile_handler_test.go @@ -223,7 +223,7 @@ func createRequest(prompt string) *scheduling.InferenceRequest { Completions: &fwkrh.CompletionsRequest{ Prompt: fwkrh.Prompt{Raw: prompt}, }, - TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{make([]uint32, len(prompt)/averageCharactersPerToken)}}, + TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: make([]uint32, len(prompt)/averageCharactersPerToken)}}}, }, } } diff --git a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/prefix_based_pd_decider.go b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/prefix_based_pd_decider.go index b2be45f5b9..0b1bdf873b 100644 --- a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/prefix_based_pd_decider.go +++ b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/prefix_based_pd_decider.go @@ -169,7 +169,7 @@ func getUserInputLenInTokens(request *scheduling.InferenceRequest) (int, error) return 0, errors.New("request or request body is nil") } - if tp := request.Body.TokenizedPrompt; tp != nil { + if tp := request.Body.TokenizedRequest; tp != nil { return tp.TokenCount(), nil } return 0, nil diff --git a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/prefix_based_pd_decider_test.go b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/prefix_based_pd_decider_test.go index fad76a2426..bbe2743f83 100644 --- a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/prefix_based_pd_decider_test.go +++ b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/prefix_based_pd_decider_test.go @@ -64,10 +64,10 @@ func makeRequestWithTokens(tokens int) *scheduling.InferenceRequest { // withTokens sets the tokenized prompt to carry n token IDs, which the decider reads // as the input token count. Any existing tokenized prompt is preserved. func withTokens(req *scheduling.InferenceRequest, n int) *scheduling.InferenceRequest { - if req.Body.TokenizedPrompt == nil { - req.Body.TokenizedPrompt = &fwkrh.TokenizedPrompt{} + if req.Body.TokenizedRequest == nil { + req.Body.TokenizedRequest = &fwkrh.TokenizedRequest{} } - req.Body.TokenizedPrompt.PerPromptTokens = [][]uint32{make([]uint32, n)} + req.Body.TokenizedRequest.Prompts = []fwkrh.PromptTokens{{TokenIDs: make([]uint32, n)}} return req } @@ -118,7 +118,7 @@ func TestGetUserInputLenInTokens(t *testing.T) { Strings: []string{"hello world", "foo bar baz"}, }, }, - TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{make([]uint32, 5)}}, + TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: make([]uint32, 5)}}}, }, }, want: 5, @@ -160,8 +160,8 @@ func TestGetUserInputLenInTokens(t *testing.T) { name: "generate request returns exact token count", req: &scheduling.InferenceRequest{ Body: &fwkrh.InferenceRequestBody{ - Generate: &fwkrh.GenerateRequest{TokenIDs: []uint32{1, 2, 3, 4, 5, 6, 7}}, - TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{make([]uint32, 7)}}, + Generate: &fwkrh.GenerateRequest{TokenIDs: []uint32{1, 2, 3, 4, 5, 6, 7}}, + TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: make([]uint32, 7)}}}, }, }, want: 7, diff --git a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/scheduler_test.go b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/scheduler_test.go index b4d78ad991..12bdfd532f 100644 --- a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/scheduler_test.go +++ b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/scheduler_test.go @@ -39,8 +39,8 @@ const ( // the input token count. func completionsBody(prompt string) *fwkrh.InferenceRequestBody { return &fwkrh.InferenceRequestBody{ - Completions: &fwkrh.CompletionsRequest{Prompt: fwkrh.Prompt{Raw: prompt}}, - TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{make([]uint32, len(prompt)/averageCharactersPerToken)}}, + Completions: &fwkrh.CompletionsRequest{Prompt: fwkrh.Prompt{Raw: prompt}}, + TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: make([]uint32, len(prompt)/averageCharactersPerToken)}}}, } } diff --git a/pkg/epp/framework/plugins/scheduling/scorer/contextlengthaware/context_length_aware.go b/pkg/epp/framework/plugins/scheduling/scorer/contextlengthaware/context_length_aware.go index 5274be05fe..09ae680906 100644 --- a/pkg/epp/framework/plugins/scheduling/scorer/contextlengthaware/context_length_aware.go +++ b/pkg/epp/framework/plugins/scheduling/scorer/contextlengthaware/context_length_aware.go @@ -80,7 +80,7 @@ func NewContextLengthAware(name string, params *contextLengthAwareParameters) *C // If filtering is enabled, endpoints that don't support the request's context length are filtered out. // Additionally, it scores endpoints based on how well their context length ranges match the request. // -// The context length is the token count from InferenceRequestBody.TokenizedPrompt as +// The context length is the token count from InferenceRequestBody.TokenizedRequest as // populated by the tokenizer DataProducer plugin. When tokens are not available it is // treated as 0 (unknown). type ContextLengthAware struct { @@ -103,12 +103,12 @@ func (p *ContextLengthAware) WithName(name string) *ContextLengthAware { return p } -// 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 plugin runs and auto-creates one when none is // configured; the context length is the token count it provides. func (p *ContextLengthAware) Consumes() plugin.DataDependencies { return plugin.DataDependencies{ - Required: map[plugin.DataKey]any{tokenproducer.TokenizedPromptDataKey: scheduling.TokenizedPrompt{}}, + Required: map[plugin.DataKey]any{tokenproducer.TokenizedPromptDataKey: scheduling.TokenizedRequest{}}, } } @@ -200,13 +200,13 @@ func (p *ContextLengthAware) Category() scheduling.ScorerCategory { } // getContextLength returns the context length (token count) for the request, read solely -// from InferenceRequestBody.TokenizedPrompt as populated by the tokenizer DataProducer +// from InferenceRequestBody.TokenizedRequest as populated by the tokenizer DataProducer // plugin. When tokens are unavailable it returns 0 (unknown). func getContextLength(request *scheduling.InferenceRequest) int { - if request == nil || request.Body == nil || request.Body.TokenizedPrompt == nil { + if request == nil || request.Body == nil || request.Body.TokenizedRequest == nil { return 0 } - return request.Body.TokenizedPrompt.TokenCount() + return request.Body.TokenizedRequest.TokenCount() } // parseContextRange parses a label value into a single context range. diff --git a/pkg/epp/framework/plugins/scheduling/scorer/contextlengthaware/context_length_aware_test.go b/pkg/epp/framework/plugins/scheduling/scorer/contextlengthaware/context_length_aware_test.go index 5c1004eb3c..8d7fd15bff 100644 --- a/pkg/epp/framework/plugins/scheduling/scorer/contextlengthaware/context_length_aware_test.go +++ b/pkg/epp/framework/plugins/scheduling/scorer/contextlengthaware/context_length_aware_test.go @@ -241,10 +241,10 @@ func TestCalculateRangeScoreFallback(t *testing.T) { }) } -// TokenizedPrompt tests — plugin reads tokens from InferenceRequestBody.TokenizedPrompt +// TokenizedRequest tests — plugin reads tokens from InferenceRequestBody.TokenizedRequest // as populated by the tokenizer DataProducer plugin. -func TestContextLengthAwareWithTokenizedPromptOnRequest(t *testing.T) { +func TestContextLengthAwareWithTokenizedRequestOnRequest(t *testing.T) { ctx := utils.NewTestContext(t) tokenCount := 42 @@ -273,7 +273,7 @@ func TestContextLengthAwareWithTokenizedPromptOnRequest(t *testing.T) { RequestID: "test-request", TargetModel: "test-model", Body: &fwkrh.InferenceRequestBody{ - TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{tokenIDs}}, + TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: tokenIDs}}}, }, } @@ -282,10 +282,10 @@ func TestContextLengthAwareWithTokenizedPromptOnRequest(t *testing.T) { assert.Equal(t, "tight-match", filteredEndpoints[0].GetMetadata().NamespacedName.Name) } -func TestContextLengthAwareNilTokenizedPromptIsZero(t *testing.T) { +func TestContextLengthAwareNilTokenizedRequestIsZero(t *testing.T) { ctx := utils.NewTestContext(t) - // Without TokenizedPrompt the context length is 0 (unknown); no protocol structs are read. + // Without TokenizedRequest the context length is 0 (unknown); no protocol structs are read. endpoints := []scheduling.Endpoint{ createEndpoint(k8stypes.NamespacedName{Namespace: "default", Name: "matching-range"}, "10.0.0.1", diff --git a/pkg/epp/framework/plugins/scheduling/scorer/preciseprefixcache/legacy_producer.go b/pkg/epp/framework/plugins/scheduling/scorer/preciseprefixcache/legacy_producer.go index e439eef2c1..eef3c1b264 100644 --- a/pkg/epp/framework/plugins/scheduling/scorer/preciseprefixcache/legacy_producer.go +++ b/pkg/epp/framework/plugins/scheduling/scorer/preciseprefixcache/legacy_producer.go @@ -84,7 +84,7 @@ func newLegacyProducer(ctx context.Context, name string, cfg preciseproducer.Plu return lp, nil } -// Consumes drops the TokenizedPrompt dependency when the wrapper owns a +// Consumes drops the TokenizedRequest dependency when the wrapper owns a // tokenizer pool, since the prompt-fallback path tokenizes the request // itself and no upstream token-producer is required. func (lp *legacyProducer) Consumes() plugin.DataDependencies { @@ -95,7 +95,7 @@ func (lp *legacyProducer) Consumes() plugin.DataDependencies { } // Produce tokenizes the request prompt via the wrapper-owned pool when -// no TokenizedPrompt is set, then delegates to the embedded Producer. +// no TokenizedRequest is set, then delegates to the embedded Producer. func (lp *legacyProducer) Produce(ctx context.Context, request *scheduling.InferenceRequest, endpoints []scheduling.Endpoint, ) error { @@ -109,7 +109,7 @@ func needsLegacyTokenization(request *scheduling.InferenceRequest) bool { if request == nil || request.Body == nil { return false } - if tp := request.Body.TokenizedPrompt; tp != nil && tp.TokenCount() > 0 { + if tp := request.Body.TokenizedRequest; tp != nil && tp.TokenCount() > 0 { return false } return request.Body.Completions != nil || request.Body.ChatCompletions != nil @@ -132,17 +132,19 @@ func (lp *legacyProducer) tokenizeRequest(request *scheduling.InferenceRequest) return } - request.Body.TokenizedPrompt = &fwkrh.TokenizedPrompt{ - PerPromptTokens: [][]uint32{tokens}, - MultiModalFeatures: flattenMMFeatures(mmFeatures), - CacheSalt: tokenizer.CacheSaltFromBody(request.Body), + request.Body.TokenizedRequest = &fwkrh.TokenizedRequest{ + Prompts: []fwkrh.PromptTokens{{ + TokenIDs: tokens, + MultiModalFeatures: flattenMMFeatures(mmFeatures), + }}, + CacheSalt: tokenizer.CacheSaltFromBody(request.Body), } } // flattenMMFeatures regroups the kvcache map-shaped multimodal metadata -// into the upstream flat list expected on TokenizedPrompt, sorted by +// into the upstream flat list expected on PromptTokens, sorted by // placeholder offset so consumers see items in prompt order. -func flattenMMFeatures(src *tokenization.MultiModalFeatures) [][]fwkrh.MultiModalFeature { +func flattenMMFeatures(src *tokenization.MultiModalFeatures) []fwkrh.MultiModalFeature { if src == nil || len(src.MMHashes) == 0 { return nil } @@ -169,5 +171,5 @@ func flattenMMFeatures(src *tokenization.MultiModalFeatures) [][]fwkrh.MultiModa return nil } sort.Slice(items, func(i, j int) bool { return items[i].Offset < items[j].Offset }) - return [][]fwkrh.MultiModalFeature{items} + return items } diff --git a/pkg/epp/framework/plugins/scheduling/scorer/preciseprefixcache/precise_prefix_cache_test.go b/pkg/epp/framework/plugins/scheduling/scorer/preciseprefixcache/precise_prefix_cache_test.go index 5d81518b66..7eda51f632 100644 --- a/pkg/epp/framework/plugins/scheduling/scorer/preciseprefixcache/precise_prefix_cache_test.go +++ b/pkg/epp/framework/plugins/scheduling/scorer/preciseprefixcache/precise_prefix_cache_test.go @@ -133,9 +133,9 @@ func (s *stubPromptTokenizer) Tokenize(rr *tokenizerTypes.RenderChatRequest, pro } // With a wrapper-owned tokenizer pool, Consumes must drop the inner -// producer's TokenizedPrompt dependency — the wrapper supplies tokens +// producer's TokenizedRequest dependency — the wrapper supplies tokens // itself and no upstream token-producer is required. -func TestLegacyProducer_ConsumesDropsTokenizedPromptWhenPoolSet(t *testing.T) { +func TestLegacyProducer_ConsumesDropsTokenizedRequestWhenPoolSet(t *testing.T) { ctx := utils.NewTestContext(t) handle := fwkplugin.NewEppHandle(ctx, nil, fwkplugin.WithMetricsRecorder(prometheus.NewRegistry())) @@ -150,12 +150,12 @@ func TestLegacyProducer_ConsumesDropsTokenizedPromptWhenPoolSet(t *testing.T) { assert.Empty(t, lp.Consumes()) lpNoPool := &legacyProducer{Producer: inner.(*preciseproducer.Producer)} - assert.NotEmpty(t, lpNoPool.Consumes(), "without a pool, Consumes must keep TokenizedPrompt") + assert.NotEmpty(t, lpNoPool.Consumes(), "without a pool, Consumes must keep TokenizedRequest") } -// When a completions prompt arrives without TokenizedPrompt and the pool +// When a completions prompt arrives without TokenizedRequest and the pool // is set, Produce must route the prompt through the pool and stash the -// resulting tokens on request.Body.TokenizedPrompt. +// resulting tokens on request.Body.TokenizedRequest. func TestLegacyProducer_TokenizesCompletionPromptViaPool(t *testing.T) { ctx := utils.NewTestContext(t) handle := fwkplugin.NewEppHandle(ctx, nil, @@ -182,11 +182,11 @@ func TestLegacyProducer_TokenizesCompletionPromptViaPool(t *testing.T) { assert.Equal(t, 1, stub.calls) assert.Equal(t, "hello world", stub.lastRaw) - require.NotNil(t, req.Body.TokenizedPrompt) - assert.Equal(t, []uint32{1, 2, 3}, req.Body.TokenizedPrompt.PerPromptTokens[0]) + require.NotNil(t, req.Body.TokenizedRequest) + assert.Equal(t, []uint32{1, 2, 3}, req.Body.TokenizedRequest.Prompts[0].TokenIDs) // Wrapper-owned tokenization must still carry the cache salt so precise // keys stay isolated on this path. - assert.Equal(t, "leg-salt", req.Body.TokenizedPrompt.CacheSalt) + assert.Equal(t, "leg-salt", req.Body.TokenizedRequest.CacheSalt) } // End-to-end: tokens from the pool flow into the embedded producer, get @@ -237,9 +237,9 @@ func TestLegacyProducer_TokensFlowToEndpointAttribute(t *testing.T) { "empty index → no matches") } -// Pre-existing TokenizedPrompt must skip the pool entirely, so the new-path +// Pre-existing TokenizedRequest must skip the pool entirely, so the new-path // token-producer pipeline isn't shadowed by the legacy pool. -func TestLegacyProducer_KeepsExistingTokenizedPrompt(t *testing.T) { +func TestLegacyProducer_KeepsExistingTokenizedRequest(t *testing.T) { ctx := utils.NewTestContext(t) handle := fwkplugin.NewEppHandle(ctx, nil, fwkplugin.WithMetricsRecorder(prometheus.NewRegistry())) @@ -255,12 +255,12 @@ func TestLegacyProducer_KeepsExistingTokenizedPrompt(t *testing.T) { req := &scheduling.InferenceRequest{ Body: &fwkrh.InferenceRequestBody{ - TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{{5, 5, 5}}}, - Completions: &fwkrh.CompletionsRequest{Prompt: fwkrh.Prompt{Raw: "should not tokenize"}}, + TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: []uint32{5, 5, 5}}}}, + Completions: &fwkrh.CompletionsRequest{Prompt: fwkrh.Prompt{Raw: "should not tokenize"}}, }, } require.NoError(t, lp.Produce(ctx, req, nil)) assert.Equal(t, 0, stub.calls, "pool must not be called when tokens already present") - assert.Equal(t, []uint32{5, 5, 5}, req.Body.TokenizedPrompt.PerPromptTokens[0]) + assert.Equal(t, []uint32{5, 5, 5}, req.Body.TokenizedRequest.Prompts[0].TokenIDs) } From 6c13146438af53cd96d257e977a27b8b0c405677 Mon Sep 17 00:00:00 2001 From: Alberto Perdomo Date: Wed, 8 Jul 2026 12:47:17 +0100 Subject: [PATCH 3/6] fix: Rename missing instances Signed-off-by: Alberto Perdomo --- .../requestcontrol/dataproducer/burstprefix/plugin.go | 2 +- .../dataproducer/burstprefix/plugin_test.go | 2 +- .../requestcontrol/dataproducer/prefixhash/hashing.go | 8 ++++---- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/plugin.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/plugin.go index 7eb3619710..6817d4197b 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/plugin.go +++ b/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{}}, } } diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/plugin_test.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/plugin_test.go index 21f641268e..0d8b0ef8a9 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/plugin_test.go +++ b/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}}}, }, } } diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/prefixhash/hashing.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/prefixhash/hashing.go index 1fbac26adc..62c1335ba8 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/prefixhash/hashing.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/prefixhash/hashing.go @@ -52,11 +52,11 @@ func (b HashBlock) Hash() uint64 { } // GetBlockHashes divides the tokenized prompt into blocks and calculates a -// prefix cache hash for each block. Each prompt in PerPromptTokens is hashed +// prefix cache hash for each block. Each prompt in Prompts is hashed // independently so cross-prompt block adjacency is avoided. The first block // hash of every prompt includes the model name and cache salt (if provided). // For subsequent blocks, the hash is calculated as: hash(block i content, hash(i-1)). -// It requires request.Body.TokenizedPrompt to be populated by a token-producer backend. +// It requires request.Body.TokenizedRequest to be populated by a token-producer backend. func GetBlockHashes(ctx context.Context, request *scheduling.InferenceRequest, blockSizeTokens int, maxPrefixBlocks int) [][]BlockHash { loggerDebug := log.FromContext(ctx).V(logutil.DEBUG) if request == nil || request.Body == nil { @@ -71,8 +71,8 @@ func GetBlockHashes(ctx context.Context, request *scheduling.InferenceRequest, b } var result [][]BlockHash - for _, tokens := range tp.PerPromptTokens { - seq := getKVCacheBlocksFromTokens(tokens, blockSizeTokens) + for _, p := range tp.Prompts { + seq := getKVCacheBlocksFromTokens(p.TokenIDs, blockSizeTokens) hashes := computeBlockHashes(seq, request, maxPrefixBlocks) if len(hashes) > 0 { result = append(result, hashes) From be02b8d05db820854cc32c4a64579886444f9b88 Mon Sep 17 00:00:00 2001 From: Alberto Perdomo Date: Mon, 13 Jul 2026 08:37:52 +0100 Subject: [PATCH 4/6] fix: Rename missing instances Signed-off-by: Alberto Perdomo --- .../requestcontrol/dataproducer/tokenizer/tokenizer_test.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/tokenizer_test.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/tokenizer_test.go index 364061278d..c059c38cc3 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/tokenizer_test.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/tokenizer_test.go @@ -792,8 +792,8 @@ func TestProduce_MessagesRequest(t *testing.T) { }, } require.NoError(t, p.Produce(context.Background(), req, nil)) - require.NotNil(t, req.Body.TokenizedPrompt) - assert.Equal(t, [][]uint32{wantTokens}, req.Body.TokenizedPrompt.PerPromptTokens) + require.NotNil(t, req.Body.TokenizedRequest) + assert.Equal(t, []fwkrh.PromptTokens{{TokenIDs: wantTokens}}, req.Body.TokenizedRequest.Prompts) pm, ok := gotPayload.AsMap() require.True(t, ok, "RenderChat payload must be a map") From 1d3ab2741d26bf7eff591d2efa88999c659b77dc Mon Sep 17 00:00:00 2001 From: Alberto Perdomo Date: Fri, 17 Jul 2026 11:19:10 +0100 Subject: [PATCH 5/6] fix: update stale TokenizedPrompt reference missed during rebase The auto-merge left one test using the old TokenizedPrompt type and PerPromptTokens field in producer_test.go. Updated to use TokenizedRequest with Prompts/PromptTokens. Signed-off-by: Alberto Perdomo --- .../dataproducer/preciseprefixcache/producer_test.go | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/producer_test.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/producer_test.go index 984aa45c2a..ccc689ef99 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/producer_test.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/producer_test.go @@ -416,8 +416,11 @@ func TestProduce_WritesCachedBlocksByTier(t *testing.T) { RequestID: "req-by-tier", TargetModel: "test-model", Body: &fwkrh.InferenceRequestBody{ - TokenizedPrompt: &fwkrh.TokenizedPrompt{ - PerPromptTokens: [][]uint32{promptA, promptB}, + TokenizedRequest: &fwkrh.TokenizedRequest{ + Prompts: []fwkrh.PromptTokens{ + {TokenIDs: promptA}, + {TokenIDs: promptB}, + }, }, }, } From ee25953be5376b83d15712798a2a3be5612e98fd Mon Sep 17 00:00:00 2001 From: Alberto Perdomo Date: Fri, 24 Jul 2026 10:44:05 +0100 Subject: [PATCH 6/6] fix: Rename stale entries Signed-off-by: Alberto Perdomo --- .../dataproducer/tokenizer/estimate_test.go | 38 +++++++++---------- 1 file changed, 19 insertions(+), 19 deletions(-) diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/estimate_test.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/estimate_test.go index 5e975cddb6..14acd02803 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/estimate_test.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/estimate_test.go @@ -235,10 +235,10 @@ func chatVideoBody(url string) *fwkrh.InferenceRequestBody { func TestVideoEstimator_Default(t *testing.T) { tp, err := estimateBackend{}.produce(context.Background(), chatVideoBody("https://example.com/clip.mp4")) require.NoError(t, err) - require.Len(t, tp.MultiModalFeatures, 1) + require.Len(t, tp.Prompts[0].MultiModalFeatures, 1) frames := defaultVideoDuration * defaultVideoSampleFPS tpf := (defaultVideoWidth * defaultVideoHeight) / videoTokenFactor - assert.Equal(t, frames*tpf, tp.MultiModalFeatures[0].Length, "default video length") + assert.Equal(t, frames*tpf, tp.Prompts[0].MultiModalFeatures[0].Length, "default video length") } // TestVideoEstimator_StaticTokensPerFrame asserts static mode emits a constant @@ -250,7 +250,7 @@ func TestVideoEstimator_StaticTokensPerFrame(t *testing.T) { tp, err := b.produce(context.Background(), chatVideoBody("https://example.com/clip.mp4")) require.NoError(t, err) frames := defaultVideoDuration * defaultVideoSampleFPS - assert.Equal(t, frames*100, tp.MultiModalFeatures[0].Length, "static tokens-per-frame video length") + assert.Equal(t, frames*100, tp.Prompts[0].MultiModalFeatures[0].Length, "static tokens-per-frame video length") } // TestVideoEstimator_DynamicFactor asserts the dynamic factor knob changes the @@ -263,7 +263,7 @@ func TestVideoEstimator_DynamicFactor(t *testing.T) { require.NoError(t, err) frames := defaultVideoDuration * defaultVideoSampleFPS tpf := (defaultVideoWidth * defaultVideoHeight) / 2048 - assert.Equal(t, frames*tpf, tp.MultiModalFeatures[0].Length, "custom-factor video length") + assert.Equal(t, frames*tpf, tp.Prompts[0].MultiModalFeatures[0].Length, "custom-factor video length") } // TestVideoEstimator_SampledFrames asserts sampled frames scale with sampleFPS @@ -276,7 +276,7 @@ func TestVideoEstimator_SampledFrames(t *testing.T) { }})} tp, err := b.produce(context.Background(), chatVideoBody("https://example.com/clip.mp4")) require.NoError(t, err) - assert.Equal(t, 8*2*10, tp.MultiModalFeatures[0].Length, "sampled-frames video length") + assert.Equal(t, 8*2*10, tp.Prompts[0].MultiModalFeatures[0].Length, "sampled-frames video length") } // TestVideoEstimator_StridedFramesCapped asserts strided frames apply the @@ -290,7 +290,7 @@ func TestVideoEstimator_StridedFramesCapped(t *testing.T) { tp, err := b.produce(context.Background(), chatVideoBody("https://example.com/clip.mp4")) require.NoError(t, err) // duration*sourceFPS/stride = 10*24/4 = 60, capped to 16; 16*100 tokens. - assert.Equal(t, 16*100, tp.MultiModalFeatures[0].Length, "strided-frames video length") + assert.Equal(t, 16*100, tp.Prompts[0].MultiModalFeatures[0].Length, "strided-frames video length") } // TestVideoEstimator_StridedFramesFloored asserts strided frames apply the @@ -304,7 +304,7 @@ func TestVideoEstimator_StridedFramesFloored(t *testing.T) { tp, err := b.produce(context.Background(), chatVideoBody("https://example.com/clip.mp4")) require.NoError(t, err) // duration*sourceFPS/stride = 1*24/4 = 6, floored to 8; 8*100 tokens. - assert.Equal(t, 8*100, tp.MultiModalFeatures[0].Length, "strided-frames floored video length") + assert.Equal(t, 8*100, tp.Prompts[0].MultiModalFeatures[0].Length, "strided-frames floored video length") } // TestVideoEstimator_MaxVideoTokens asserts the overall cap bounds the total @@ -316,7 +316,7 @@ func TestVideoEstimator_MaxVideoTokens(t *testing.T) { tp, err := b.produce(context.Background(), chatVideoBody("https://example.com/clip.mp4")) require.NoError(t, err) // Default frames*tpf = 10*225 = 2250, capped to 500. - assert.Equal(t, 500, tp.MultiModalFeatures[0].Length, "max-video-tokens cap") + assert.Equal(t, 500, tp.Prompts[0].MultiModalFeatures[0].Length, "max-video-tokens cap") } // TestVideoEstimator_Qwen3AndGemma4 asserts the two documented model shapes @@ -332,7 +332,7 @@ func TestVideoEstimator_Qwen3AndGemma4(t *testing.T) { tp, err := qwen3.produce(context.Background(), chatVideoBody("https://example.com/clip.mp4")) require.NoError(t, err) // frames = 10*2 = 20, tpf = 640*480/1024 = 300, tokens = 6000. - assert.Equal(t, 20*((640*480)/1024), tp.MultiModalFeatures[0].Length, "qwen3-shaped video length") + assert.Equal(t, 20*((640*480)/1024), tp.Prompts[0].MultiModalFeatures[0].Length, "qwen3-shaped video length") gemma4 := estimateBackend{vid: newVideoEstimator(&estimateConfig{Video: &videoEstimateConfig{ DefaultDuration: 10, @@ -342,7 +342,7 @@ func TestVideoEstimator_Qwen3AndGemma4(t *testing.T) { tp, err = gemma4.produce(context.Background(), chatVideoBody("https://example.com/clip.mp4")) require.NoError(t, err) // frames = min(10*24/4, 16) = 16, tokens = 16*256. - assert.Equal(t, 16*256, tp.MultiModalFeatures[0].Length, "gemma4-shaped video length") + assert.Equal(t, 16*256, tp.Prompts[0].MultiModalFeatures[0].Length, "gemma4-shaped video length") } // TestParseVideoMetadataHeaders covers full, partial, missing, and malformed @@ -433,11 +433,11 @@ func TestVideoEstimator_HeaderMetadataOverridesDefaults(t *testing.T) { withMeta, err := estimateBackend{}.produce(ctx, chatVideoBody("https://example.com/clip.mp4")) require.NoError(t, err) // sampled frames = duration(3)*sampleFPS(2) = 6; dynamic tpf = 320*240/1024 = 75. - assert.Equal(t, 6*((320*240)/videoTokenFactor), withMeta.MultiModalFeatures[0].Length) + assert.Equal(t, 6*((320*240)/videoTokenFactor), withMeta.Prompts[0].MultiModalFeatures[0].Length) def, err := estimateBackend{}.produce(context.Background(), chatVideoBody("https://example.com/clip.mp4")) require.NoError(t, err) - assert.NotEqual(t, def.MultiModalFeatures[0].Length, withMeta.MultiModalFeatures[0].Length, "header metadata must change the count") + assert.NotEqual(t, def.Prompts[0].MultiModalFeatures[0].Length, withMeta.Prompts[0].MultiModalFeatures[0].Length, "header metadata must change the count") } // TestVideoEstimator_HeaderFPSStridedMode asserts header source FPS and duration @@ -450,7 +450,7 @@ func TestVideoEstimator_HeaderFPSStridedMode(t *testing.T) { // strided frames = int(duration(3)*fps(30))/2 = 45. tp, err := b.produce(videoCtx(videoMetadata{duration: 3, fps: 30}), chatVideoBody("https://cdn.example.com/movie.mp4")) require.NoError(t, err) - assert.Equal(t, 45, tp.MultiModalFeatures[0].Length) + assert.Equal(t, 45, tp.Prompts[0].MultiModalFeatures[0].Length) } // TestVideoEstimator_SampledIgnoresHeaderFPS asserts sampled mode honors the @@ -465,8 +465,8 @@ func TestVideoEstimator_SampledIgnoresHeaderFPS(t *testing.T) { require.NoError(t, err) fps60, err := b.produce(videoCtx(videoMetadata{duration: 3, fps: 60}), chatVideoBody("https://example.com/clip.mp4")) require.NoError(t, err) - assert.Equal(t, 6, fps30.MultiModalFeatures[0].Length) - assert.Equal(t, fps30.MultiModalFeatures[0].Length, fps60.MultiModalFeatures[0].Length, "sampled mode must ignore header source fps") + assert.Equal(t, 6, fps30.Prompts[0].MultiModalFeatures[0].Length) + assert.Equal(t, fps30.Prompts[0].MultiModalFeatures[0].Length, fps60.Prompts[0].MultiModalFeatures[0].Length, "sampled mode must ignore header source fps") } // TestVideoEstimator_NoHeadersUseDefaults asserts that without header metadata the @@ -476,7 +476,7 @@ func TestVideoEstimator_NoHeadersUseDefaults(t *testing.T) { require.NoError(t, err) frames := defaultVideoDuration * defaultVideoSampleFPS tpf := (defaultVideoWidth * defaultVideoHeight) / videoTokenFactor - assert.Equal(t, frames*tpf, tp.MultiModalFeatures[0].Length, "default video length") + assert.Equal(t, frames*tpf, tp.Prompts[0].MultiModalFeatures[0].Length, "default video length") } // TestVideoEstimator_TemporalMergeAndMinFrames asserts sampled frames are floored @@ -490,11 +490,11 @@ func TestVideoEstimator_TemporalMergeAndMinFrames(t *testing.T) { // 10s: 10*2 = 20 sampled frames, /2 temporal merge = 10 groups. long, err := b.produce(videoCtx(videoMetadata{duration: 10}), chatVideoBody("https://example.com/clip.mp4")) require.NoError(t, err) - assert.Equal(t, 10, long.MultiModalFeatures[0].Length) + assert.Equal(t, 10, long.Prompts[0].MultiModalFeatures[0].Length) // 1s: 2 sampled frames floored to minFrames 4, /2 merge = 2 groups. short, err := b.produce(videoCtx(videoMetadata{duration: 1}), chatVideoBody("https://example.com/clip.mp4")) require.NoError(t, err) - assert.Equal(t, 2, short.MultiModalFeatures[0].Length) + assert.Equal(t, 2, short.Prompts[0].MultiModalFeatures[0].Length) } // TestVideoEstimator_HeaderRespectsMaxVideoTokens asserts the overall cap still @@ -504,7 +504,7 @@ func TestVideoEstimator_HeaderRespectsMaxVideoTokens(t *testing.T) { // Uncapped would be 6*75 = 450; capped to 100. tp, err := b.produce(videoCtx(videoMetadata{width: 320, height: 240, duration: 3}), chatVideoBody("https://example.com/clip.mp4")) require.NoError(t, err) - assert.Equal(t, 100, tp.MultiModalFeatures[0].Length) + assert.Equal(t, 100, tp.Prompts[0].MultiModalFeatures[0].Length) } // TestEstimateBackend_MessagesImageFeature asserts an Anthropic messages image