Skip to content

Commit e781181

Browse files
albertoperdomo2albertoperdomo2
committed
fix: Bundle per-prompt tokens and MM features into PromptTokens struct
Signed-off-by: Alberto Perdomo <aperdomo@redhat.com>
1 parent 27a62e4 commit e781181

43 files changed

Lines changed: 370 additions & 352 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎pkg/epp/flowcontrol/benchmark/benchmark_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -435,7 +435,7 @@ func BenchmarkFlowController_FullPathStress(b *testing.B) {
435435
// 2. Post-scheduling: producer tracks the request on the endpoint.
436436
infReq := &scheduling.InferenceRequest{
437437
RequestID: reqID,
438-
Body: &requesthandling.InferenceRequestBody{TokenizedPrompt: &requesthandling.TokenizedPrompt{PerPromptTokens: [][]uint32{benchTokenIDs}}},
438+
Body: &requesthandling.InferenceRequestBody{TokenizedRequest: &requesthandling.TokenizedRequest{Prompts: []requesthandling.PromptTokens{{TokenIDs: benchTokenIDs}}}},
439439
}
440440
schedResult := &scheduling.SchedulingResult{ProfileResults: profileResults}
441441
h.producer.PreRequest(ctx, infReq, schedResult)

‎pkg/epp/flowcontrol/integration_test.go‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,7 @@ func TestConcurrentSaturationReads(t *testing.T) {
8282
req := &fwksched.InferenceRequest{
8383
RequestID: fmt.Sprintf("req-%d", i),
8484
Body: &fwkrh.InferenceRequestBody{
85-
TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{make([]uint32, 10)}},
85+
TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: make([]uint32, 10)}}},
8686
},
8787
}
8888
result := &fwksched.SchedulingResult{
@@ -155,7 +155,7 @@ func TestSaturationFullLoop(t *testing.T) {
155155
req := &fwksched.InferenceRequest{
156156
RequestID: fmt.Sprintf("prefill-%d", i),
157157
Body: &fwkrh.InferenceRequestBody{
158-
TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{make([]uint32, 50)}},
158+
TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: make([]uint32, 50)}}},
159159
},
160160
}
161161
result := &fwksched.SchedulingResult{
@@ -444,7 +444,7 @@ func TestUsageLimitThresholdGatesDispatch(t *testing.T) {
444444
req := &fwksched.InferenceRequest{
445445
RequestID: fmt.Sprintf("inflight-%d", i),
446446
Body: &fwkrh.InferenceRequestBody{
447-
TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{make([]uint32, 10)}},
447+
TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: make([]uint32, 10)}}},
448448
},
449449
}
450450
result := &fwksched.SchedulingResult{
@@ -1005,7 +1005,7 @@ func TestEndpointReregistrationSaturationAccuracy(t *testing.T) {
10051005
oldReq := &fwksched.InferenceRequest{
10061006
RequestID: "old-req",
10071007
Body: &fwkrh.InferenceRequestBody{
1008-
TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{make([]uint32, 50)}},
1008+
TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: make([]uint32, 50)}}},
10091009
},
10101010
}
10111011
oldResult := &fwksched.SchedulingResult{
@@ -1061,7 +1061,7 @@ func TestEndpointReregistrationSaturationAccuracy(t *testing.T) {
10611061
newReq := &fwksched.InferenceRequest{
10621062
RequestID: "new-req",
10631063
Body: &fwkrh.InferenceRequestBody{
1064-
TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{make([]uint32, 50)}},
1064+
TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: make([]uint32, 50)}}},
10651065
},
10661066
}
10671067
newResult := &fwksched.SchedulingResult{
@@ -1122,7 +1122,7 @@ func TestEndpointIdentityCollisionDuringPodReplacement(t *testing.T) {
11221122
req := &fwksched.InferenceRequest{
11231123
RequestID: "new-pod-req",
11241124
Body: &fwkrh.InferenceRequestBody{
1125-
TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{make([]uint32, 50)}},
1125+
TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: make([]uint32, 50)}}},
11261126
},
11271127
}
11281128
result := &fwksched.SchedulingResult{

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

Lines changed: 22 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -95,9 +95,9 @@ type InferenceRequestBody struct {
9595
// If the payload is unmarshaled, we can perform advanced processing (like prefix cache aware routing).
9696
// If it remains as raw bytes, such processing may not be supported.
9797
Payload RequestPayload `json:"-"`
98-
// TokenizedPrompt contains parser-derived tokenization results when available.
98+
// TokenizedRequest contains parser-derived tokenization results when available.
9999
// It is nil when the request was not already tokenized.
100-
TokenizedPrompt *TokenizedPrompt `json:"-"`
100+
TokenizedRequest *TokenizedRequest `json:"-"`
101101

102102
// Stream indicates whether the request specifies a streaming response (e.g., via a stream field).
103103
// This typically implies the model server's response will be streamed.
@@ -148,31 +148,36 @@ func MaxOutputTokensFromPayload(m PayloadMap, keys ...string) *int64 {
148148
return nil
149149
}
150150

151-
// TokenizedPrompt contains the result of tokenizing the request prompt.
151+
// TokenizedRequest contains the result of tokenizing the request prompt.
152152
// It is consumed by scheduling and request-control plugins that benefit from
153153
// actual token data such as prefix-cache awareness.
154-
type TokenizedPrompt struct {
155-
// PerPromptTokens holds the token IDs for each prompt in the request.
156-
// Single-prompt requests (chat, generate, single-string completions) use a
157-
// length-1 outer slice. Multi-string completions use one inner slice per
158-
// prompt string.
159-
PerPromptTokens [][]uint32
160-
// MultiModalFeatures holds multimodal items per prompt, indexed in
161-
// lockstep with PerPromptTokens. Single-prompt requests use a length-1
162-
// outer slice. Nil if the prompt contains no multimodal content.
163-
MultiModalFeatures [][]MultiModalFeature
154+
type TokenizedRequest struct {
155+
// Prompts holds the per-prompt token data. Single-prompt requests (chat,
156+
// generate, single-string completions) use a length-1 slice. Multi-string
157+
// completions use one entry per prompt string.
158+
Prompts []PromptTokens
164159
// CacheSalt isolates prefix caches across requests. Populated by the token-producer.
165160
CacheSalt string
166161
}
167162

163+
// PromptTokens bundles the token IDs and multimodal features for a single
164+
// prompt in the request.
165+
type PromptTokens struct {
166+
// TokenIDs holds the token IDs for this prompt.
167+
TokenIDs []uint32
168+
// MultiModalFeatures holds multimodal items for this prompt, ordered by
169+
// token position. Nil if the prompt contains no multimodal content.
170+
MultiModalFeatures []MultiModalFeature
171+
}
172+
168173
// TokenCount returns the total number of tokens across all prompts.
169-
func (tp *TokenizedPrompt) TokenCount() int {
174+
func (tp *TokenizedRequest) TokenCount() int {
170175
if tp == nil {
171176
return 0
172177
}
173178
n := 0
174-
for _, pp := range tp.PerPromptTokens {
175-
n += len(pp)
179+
for _, p := range tp.Prompts {
180+
n += len(p.TokenIDs)
176181
}
177182
return n
178183
}
@@ -186,7 +191,7 @@ type MultiModalFeature struct {
186191
// Hash is the content hash of the item, used for KV-cache reuse across requests.
187192
Hash string
188193
// Offset is the index of the first placeholder token for this item
189-
// in the corresponding PerPromptTokens entry.
194+
// in the owning PromptTokens.TokenIDs slice.
190195
Offset int
191196
// Length is the number of placeholder tokens this item occupies.
192197
Length int

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ type Modality = fwkrh.Modality
3232

3333
const ModalityImage = fwkrh.ModalityImage
3434

35-
type TokenizedPrompt = fwkrh.TokenizedPrompt
35+
type TokenizedRequest = fwkrh.TokenizedRequest
3636

3737
type MultiModalFeature = fwkrh.MultiModalFeature
3838

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -697,7 +697,7 @@ func makeTokenRequest(requestID string, inputTokens int) *fwksched.InferenceRequ
697697
return &fwksched.InferenceRequest{
698698
RequestID: requestID,
699699
Body: &fwkrh.InferenceRequestBody{
700-
TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{make([]uint32, inputTokens)}},
700+
TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: make([]uint32, inputTokens)}}},
701701
},
702702
}
703703
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -75,12 +75,12 @@ func (p *dataProducer) Produces() map[plugin.DataKey]any {
7575
return map[plugin.DataKey]any{p.dk: attrprefix.PrefixCacheMatchInfo{}}
7676
}
7777

78-
// Consumes declares the TokenizedPrompt dependency so the data-layer DAG orders
78+
// Consumes declares the TokenizedRequest dependency so the data-layer DAG orders
7979
// the token-producer before this producer runs and auto-creates one when none
8080
// is configured.
8181
func (p *dataProducer) Consumes() plugin.DataDependencies {
8282
return plugin.DataDependencies{
83-
Required: map[plugin.DataKey]any{tokenproducer.TokenizedPromptDataKey: fwksched.TokenizedPrompt{}},
83+
Required: map[plugin.DataKey]any{tokenproducer.TokenizedPromptDataKey: fwksched.TokenizedRequest{}},
8484
}
8585
}
8686

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

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ func disableMinBlockSizeClamp(t *testing.T) {
5353
// tokenizedBody returns a request body carrying only a tokenized prompt.
5454
func tokenizedBody(tokenIDs []uint32) *fwkrh.InferenceRequestBody {
5555
return &fwkrh.InferenceRequestBody{
56-
TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{tokenIDs}},
56+
TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: tokenIDs}}},
5757
}
5858
}
5959

@@ -629,8 +629,8 @@ func TestProduce_MultiPrompt(t *testing.T) {
629629
RequestID: uuid.NewString(),
630630
TargetModel: "test-model",
631631
Body: &fwkrh.InferenceRequestBody{
632-
TokenizedPrompt: &fwkrh.TokenizedPrompt{
633-
PerPromptTokens: [][]uint32{{1, 2, 3}, {4, 5}},
632+
TokenizedRequest: &fwkrh.TokenizedRequest{
633+
Prompts: []fwkrh.PromptTokens{{TokenIDs: []uint32{1, 2, 3}}, {TokenIDs: []uint32{4, 5}}},
634634
},
635635
},
636636
}
@@ -672,8 +672,8 @@ func TestMultiPromptMatchAggregation(t *testing.T) {
672672
RequestID: uuid.NewString(),
673673
TargetModel: "test-model",
674674
Body: &fwkrh.InferenceRequestBody{
675-
TokenizedPrompt: &fwkrh.TokenizedPrompt{
676-
PerPromptTokens: [][]uint32{{1, 2, 3}, {4, 5}},
675+
TokenizedRequest: &fwkrh.TokenizedRequest{
676+
Prompts: []fwkrh.PromptTokens{{TokenIDs: []uint32{1, 2, 3}}, {TokenIDs: []uint32{4, 5}}},
677677
},
678678
},
679679
}
@@ -691,8 +691,8 @@ func TestMultiPromptMatchAggregation(t *testing.T) {
691691
RequestID: uuid.NewString(),
692692
TargetModel: "test-model",
693693
Body: &fwkrh.InferenceRequestBody{
694-
TokenizedPrompt: &fwkrh.TokenizedPrompt{
695-
PerPromptTokens: [][]uint32{{1, 2, 3}, {4, 5}},
694+
TokenizedRequest: &fwkrh.TokenizedRequest{
695+
Prompts: []fwkrh.PromptTokens{{TokenIDs: []uint32{1, 2, 3}}, {TokenIDs: []uint32{4, 5}}},
696696
},
697697
},
698698
}
@@ -725,8 +725,8 @@ func TestMultiPromptPartialMatch(t *testing.T) {
725725
RequestID: uuid.NewString(),
726726
TargetModel: "test-model",
727727
Body: &fwkrh.InferenceRequestBody{
728-
TokenizedPrompt: &fwkrh.TokenizedPrompt{
729-
PerPromptTokens: [][]uint32{{1, 2}, {3, 4}},
728+
TokenizedRequest: &fwkrh.TokenizedRequest{
729+
Prompts: []fwkrh.PromptTokens{{TokenIDs: []uint32{1, 2}}, {TokenIDs: []uint32{3, 4}}},
730730
},
731731
},
732732
}
@@ -744,8 +744,8 @@ func TestMultiPromptPartialMatch(t *testing.T) {
744744
RequestID: uuid.NewString(),
745745
TargetModel: "test-model",
746746
Body: &fwkrh.InferenceRequestBody{
747-
TokenizedPrompt: &fwkrh.TokenizedPrompt{
748-
PerPromptTokens: [][]uint32{{1, 2}, {5, 6}},
747+
TokenizedRequest: &fwkrh.TokenizedRequest{
748+
Prompts: []fwkrh.PromptTokens{{TokenIDs: []uint32{1, 2}}, {TokenIDs: []uint32{5, 6}}},
749749
},
750750
},
751751
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -563,7 +563,7 @@ func (p *InFlightLoadProducer) Produces() map[fwkplugin.DataKey]any {
563563
}
564564
}
565565

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

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ func TestInFlightLoadProducer_Consumes(t *testing.T) {
5757

5858
deps := newTestProducer(t).Consumes()
5959

60-
// TokenizedPrompt is required so the data-layer DAG auto-creates a
60+
// TokenizedRequest is required so the data-layer DAG auto-creates a
6161
// token-producer and orders it ahead of this producer; without it the input
6262
// token estimate silently reads zero.
6363
require.Contains(t, deps.Required, tokenproducer.TokenizedPromptDataKey)
@@ -587,8 +587,8 @@ func makeTokenRequest(requestID string, inputTokens int) *fwksched.InferenceRequ
587587
return &fwksched.InferenceRequest{
588588
RequestID: requestID,
589589
Body: &fwkrh.InferenceRequestBody{
590-
TokenizedPrompt: &fwkrh.TokenizedPrompt{
591-
PerPromptTokens: [][]uint32{make([]uint32, inputTokens)},
590+
TokenizedRequest: &fwkrh.TokenizedRequest{
591+
Prompts: []fwkrh.PromptTokens{{TokenIDs: make([]uint32, inputTokens)}},
592592
},
593593
},
594594
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -87,10 +87,10 @@ func (e *SimpleTokenEstimator) Estimate(request *fwksched.InferenceRequest) int6
8787
// EstimateInput returns the input token count read from the tokenized prompt,
8888
// or 0 when no tokenization is available.
8989
func (e *SimpleTokenEstimator) EstimateInput(request *fwksched.InferenceRequest) int64 {
90-
if request == nil || request.Body == nil || request.Body.TokenizedPrompt == nil {
90+
if request == nil || request.Body == nil || request.Body.TokenizedRequest == nil {
9191
return 0
9292
}
93-
return int64(request.Body.TokenizedPrompt.TokenCount())
93+
return int64(request.Body.TokenizedRequest.TokenCount())
9494
}
9595

9696
// EstimateOutput returns the estimated output token count given the input token

0 commit comments

Comments
 (0)