Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@ DataProducer, PreRequest, ResponseHeader, ResponseBody, ProducerPlugin, Consumer
- TPOT training data collection at EOS (streaming mode)
- Per-endpoint running request queue tracking (TPOT SLO priority queue)
- Prefix cache score forwarding from `PrefixCacheMatchInfo` attributes
- Multimodal encoder-cache size forwarding from `EncoderCacheMatchInfo` attributes
(opt-in via `useEncoderCacheFeatures`)
- TPOT neutralization for prefill endpoints in disaggregated serving
- E2E latency metrics when `streamingMode=false`

Expand All @@ -30,6 +32,8 @@ DataProducer, PreRequest, ResponseHeader, ResponseBody, ProducerPlugin, Consumer
| `streamingMode` | `false` | Record TTFT on first chunk (true) vs EOS (false) |
| `endpointRoleLabel` | `""` | Label key for disaggregated serving roles |
| `predictInProduce` | `true` | Enable/disable bulk predictions. Set false for training-only mode |
| `useEncoderCacheFeatures` | `false` | Feed multimodal encoder-cache sizes (`encoder_input_size`, `encoder_matched_size`) to the predictor. Requires (and auto-creates) a multimodal encoder-cache producer |
| `encoderCacheMatchInfoProducerName` | `""` | Multimodal encoder-cache producer to read match data from. Empty defaults to the auto-created producer |

## Default Behavior (`streamingMode: false`)

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ import (
fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling"
attrconcurrency "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/concurrency"
attrlatency "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/latency"
attrmm "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/multimodal"
attrprefix "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/prefix"
tokenproducer "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer"
)
Expand Down Expand Up @@ -59,6 +60,10 @@ func (pl *PredictedLatency) Produce(ctx context.Context, request *fwksched.Infer
}
predictedLatencyCtx.prefixCacheScoresForEndpoints[endpoint.GetMetadata().NamespacedName.Name] = prefixCacheScore

if pl.config.UseEncoderCacheFeatures {
pl.captureEncoderCacheSizes(ctx, predictedLatencyCtx, endpoint)
}

// Capture the in-flight load here, in the DAG-ordered Produce hook, and
// reuse it for both the prediction features and the dispatch-time training
// features. Re-reading the live attribute in PreRequest would make the
Expand Down Expand Up @@ -114,18 +119,64 @@ func (pl *PredictedLatency) Produce(ctx context.Context, request *fwksched.Infer
return nil
}

// captureEncoderCacheSizes stores the request's total multimodal encoder item
// size and the endpoint's matched portion in the request context. The match
// data is attached by the multimodal encoder-cache producer, which the DAG
// orders before this plugin; absence means a text-only request and leaves the
// sizes at 0.
func (pl *PredictedLatency) captureEncoderCacheSizes(ctx context.Context, predictedLatencyCtx *predictedLatencyCtx, endpoint fwksched.Endpoint) {
logger := log.FromContext(ctx)
raw, ok := endpoint.Get(pl.encoderCacheDataKey.String())
if !ok {
return
}
matchInfo, ok := raw.(*attrmm.EncoderCacheMatchInfo)
if !ok || matchInfo == nil {
return
}
inputSize := sumItemSizes(matchInfo.RequestItems())
matchedSize := sumItemSizes(matchInfo.MatchedItems())
// The predictor rejects matched > input, so clamp inconsistent match data
// rather than dropping the whole prediction or training entry.
if matchedSize > inputSize {
logger.V(logutil.DEBUG).Info("Encoder cache matched size exceeds input size, clamping",
"pod", endpoint.GetMetadata().String(),
"encoderInputSize", inputSize,
"encoderMatchedSize", matchedSize)
matchedSize = inputSize
}
predictedLatencyCtx.encoderInputSize = inputSize
predictedLatencyCtx.encoderMatchedSizeForEndpoints[endpoint.GetMetadata().NamespacedName.Name] = matchedSize
logger.V(logutil.DEBUG).Info("Encoder cache sizes for pod",
"pod", endpoint.GetMetadata().String(),
"encoderInputSize", inputSize,
"encoderMatchedSize", matchedSize)
}

func sumItemSizes(items []attrmm.MatchItem) int {
total := 0
for _, item := range items {
total += item.Size
}
return total
}

func (pl *PredictedLatency) Produces() map[plugin.DataKey]any {
return map[plugin.DataKey]any{
pl.latencyPredictionInfoDataKey: attrlatency.LatencyPredictionInfo{},
}
}

func (pl *PredictedLatency) Consumes() plugin.DataDependencies {
return plugin.DataDependencies{
Required: map[plugin.DataKey]any{
pl.prefixMatchDataKey: attrprefix.PrefixCacheMatchInfo{},
pl.inFlightLoadDataKey: attrconcurrency.InFlightLoad{},
tokenproducer.TokenizedPromptDataKey: fwksched.TokenizedPrompt{},
},
required := map[plugin.DataKey]any{
pl.prefixMatchDataKey: attrprefix.PrefixCacheMatchInfo{},
pl.inFlightLoadDataKey: attrconcurrency.InFlightLoad{},
tokenproducer.TokenizedPromptDataKey: fwksched.TokenizedPrompt{},
}
// Required (not Optional) because only Required dependencies create DAG
// ordering edges; the encoder-cache producer must run before this plugin.
if pl.config.UseEncoderCacheFeatures {
required[pl.encoderCacheDataKey] = attrmm.EncoderCacheMatchInfo{}
}
return plugin.DataDependencies{Required: required}
}
Original file line number Diff line number Diff line change
Expand Up @@ -21,9 +21,11 @@ import (
"testing"

"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"

fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling"
attrlatency "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/latency"
attrmm "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/multimodal"
attrprefix "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/prefix"
)

Expand All @@ -36,6 +38,88 @@ func TestProducesConsumes(t *testing.T) {

consumes := pl.Consumes()
assert.Contains(t, consumes.Required, attrprefix.PrefixCacheMatchInfoDataKey)
assert.NotContains(t, consumes.Required, attrmm.EncoderCacheMatchInfoKey,
"encoder-cache match data must not be consumed when the feature is disabled")
}

func TestConsumes_EncoderCacheFeatureEnabled(t *testing.T) {
cfg := DefaultConfig
cfg.UseEncoderCacheFeatures = true
pl := NewPredictedLatency(LatencyDataProviderPluginType, cfg, nil)

consumes := pl.Consumes()
assert.Contains(t, consumes.Required, attrmm.EncoderCacheMatchInfoKey)
}

// TestProduce_CapturesEncoderCacheSizes verifies that Produce reads the
// multimodal encoder-cache match data attached to endpoints and captures the
// request's input size plus the per-endpoint matched size, leaving endpoints
// without match data (text-only requests) at 0.
func TestProduce_CapturesEncoderCacheSizes(t *testing.T) {
cfg := DefaultConfig
cfg.PredictInProduce = false
cfg.UseEncoderCacheFeatures = true
pl := NewPredictedLatency(LatencyDataProviderPluginType, cfg, nil)

request := createTestInferenceRequest("encoder-test", 0, 0)
matched := createTestEndpoint("pod-matched", 0.1, 0, 0)
unmatched := createTestEndpoint("pod-unmatched", 0.1, 0, 0)

items := []attrmm.MatchItem{{Hash: "img-a", Size: 1}, {Hash: "img-b", Size: 1}}
matched.Put(pl.encoderCacheDataKey.String(), attrmm.NewEncoderCacheMatchInfo(items[:1], items))
unmatched.Put(pl.encoderCacheDataKey.String(), attrmm.NewEncoderCacheMatchInfo(nil, items))

require.NoError(t, pl.Produce(context.Background(), request, []fwksched.Endpoint{matched, unmatched}))

plCtx, err := pl.getPredictedLatencyContextForRequest(request)
require.NoError(t, err)
assert.Equal(t, 2, plCtx.encoderInputSize)
assert.Equal(t, 1, plCtx.encoderMatchedSizeForEndpoints["pod-matched"])
assert.Equal(t, 0, plCtx.encoderMatchedSizeForEndpoints["pod-unmatched"])
}

// TestProduce_ClampsInconsistentEncoderMatchData verifies that match data
// whose matched size exceeds the input size is clamped so downstream
// predictor validation (matched <= input) cannot reject the request.
func TestProduce_ClampsInconsistentEncoderMatchData(t *testing.T) {
cfg := DefaultConfig
cfg.PredictInProduce = false
cfg.UseEncoderCacheFeatures = true
pl := NewPredictedLatency(LatencyDataProviderPluginType, cfg, nil)

request := createTestInferenceRequest("encoder-clamp-test", 0, 0)
endpoint := createTestEndpoint("pod-a", 0.1, 0, 0)

matched := []attrmm.MatchItem{{Hash: "img-a", Size: 3}}
requestItems := []attrmm.MatchItem{{Hash: "img-b", Size: 1}}
endpoint.Put(pl.encoderCacheDataKey.String(), attrmm.NewEncoderCacheMatchInfo(matched, requestItems))

require.NoError(t, pl.Produce(context.Background(), request, []fwksched.Endpoint{endpoint}))

plCtx, err := pl.getPredictedLatencyContextForRequest(request)
require.NoError(t, err)
assert.Equal(t, 1, plCtx.encoderInputSize)
assert.Equal(t, 1, plCtx.encoderMatchedSizeForEndpoints["pod-a"])
}

// TestProduce_EncoderCacheFeatureDisabledIgnoresMatchData is the negative
// control: with the feature off, attached match data is not read.
func TestProduce_EncoderCacheFeatureDisabledIgnoresMatchData(t *testing.T) {
cfg := DefaultConfig
cfg.PredictInProduce = false
pl := NewPredictedLatency(LatencyDataProviderPluginType, cfg, nil)

request := createTestInferenceRequest("encoder-disabled-test", 0, 0)
endpoint := createTestEndpoint("pod-a", 0.1, 0, 0)
items := []attrmm.MatchItem{{Hash: "img-a", Size: 1}}
endpoint.Put(pl.encoderCacheDataKey.String(), attrmm.NewEncoderCacheMatchInfo(items, items))

require.NoError(t, pl.Produce(context.Background(), request, []fwksched.Endpoint{endpoint}))

plCtx, err := pl.getPredictedLatencyContextForRequest(request)
require.NoError(t, err)
assert.Equal(t, 0, plCtx.encoderInputSize)
assert.Empty(t, plCtx.encoderMatchedSizeForEndpoints)
}

// TestProduce_CancelledContextDoesNotPublish verifies that when the
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -167,13 +167,16 @@ func (p *Predictor) predictBayesianRidge(req PredictionRequest, mr *MetricsRespo
}
c := mr.Coefficients

// Updated linear combination for TTFT to include prefix_cache_score
// Linear combination for TTFT; encoder-cache coefficients are absent from
// models trained without multimodal features and contribute 0 in that case.
ttft := c.TTFTIntercept +
c.TTFTCoeffs["kv_cache_percentage"]*req.KVCachePercentage +
c.TTFTCoeffs["input_token_length"]*float64(req.InputTokenLength) +
c.TTFTCoeffs["num_request_waiting"]*float64(req.NumRequestWaiting) +
c.TTFTCoeffs["num_request_running"]*float64(req.NumRequestRunning) +
c.TTFTCoeffs["prefix_cache_score"]*req.PrefixCacheScore
c.TTFTCoeffs["prefix_cache_score"]*req.PrefixCacheScore +
c.TTFTCoeffs["encoder_input_size"]*float64(req.EncoderInputSize) +
c.TTFTCoeffs["encoder_matched_size"]*float64(req.EncoderMatchedSize)

// Linear combination for TPOT (remains unchanged - no prefix cache effect)
tpot := c.TPOTIntercept +
Expand Down Expand Up @@ -252,6 +255,12 @@ func (p *Predictor) ValidatePredictionRequest(req PredictionRequest) error {
if req.PrefixCacheScore < 0.0 || req.PrefixCacheScore > 1.0 {
return fmt.Errorf("prefix_cache_score must be between 0.0 and 1.0, got %f", req.PrefixCacheScore)
}
if req.EncoderInputSize < 0 {
return fmt.Errorf("encoder_input_size must be non-negative, got %d", req.EncoderInputSize)
}
if req.EncoderMatchedSize < 0 || req.EncoderMatchedSize > req.EncoderInputSize {
return fmt.Errorf("encoder_matched_size must be between 0 and encoder_input_size (%d), got %d", req.EncoderInputSize, req.EncoderMatchedSize)
}
return nil
}

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,118 @@
/*
Copyright 2026 The llm-d Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package latencypredictorclient

import (
"testing"

"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)

func validEncoderPredictionRequest() PredictionRequest {
return PredictionRequest{
KVCachePercentage: 0.5,
InputTokenLength: 100,
NumRequestWaiting: 1,
NumRequestRunning: 1,
NumTokensGenerated: 10,
}
}

func TestValidatePredictionRequest_EncoderSizes(t *testing.T) {
predictor := &Predictor{}

tests := []struct {
name string
inputSize int
matchedSize int
wantErr bool
}{
{name: "both zero", inputSize: 0, matchedSize: 0, wantErr: false},
{name: "partial match", inputSize: 5, matchedSize: 3, wantErr: false},
{name: "full match", inputSize: 5, matchedSize: 5, wantErr: false},
{name: "negative input size", inputSize: -1, matchedSize: 0, wantErr: true},
{name: "negative matched size", inputSize: 5, matchedSize: -1, wantErr: true},
{name: "matched exceeds input", inputSize: 5, matchedSize: 6, wantErr: true},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
req := validEncoderPredictionRequest()
req.EncoderInputSize = tt.inputSize
req.EncoderMatchedSize = tt.matchedSize
err := predictor.ValidatePredictionRequest(req)
if tt.wantErr {
assert.Error(t, err)
} else {
assert.NoError(t, err)
}
})
}
}

func TestValidateTrainingEntry_EncoderSizes(t *testing.T) {
predictor := &Predictor{}
entry := TrainingEntry{
KVCachePercentage: 0.5,
InputTokenLength: 100,
NumRequestWaiting: 1,
NumRequestRunning: 1,
NumTokensGenerated: 10,
ActualTTFT: 50.0,
ActualTPOT: 5.0,
}

entry.EncoderInputSize, entry.EncoderMatchedSize = 4, 2
assert.NoError(t, predictor.ValidateTrainingEntry(entry))

entry.EncoderInputSize, entry.EncoderMatchedSize = -1, 0
assert.Error(t, predictor.ValidateTrainingEntry(entry))

entry.EncoderInputSize, entry.EncoderMatchedSize = 2, 3
assert.Error(t, predictor.ValidateTrainingEntry(entry))
}

// TestPredictBayesianRidge_EncoderCoefficients verifies the TTFT linear model
// applies encoder-cache coefficients when present and degrades to a zero
// contribution when the model was trained without them.
func TestPredictBayesianRidge_EncoderCoefficients(t *testing.T) {
predictor := &Predictor{}
req := validEncoderPredictionRequest()
req.EncoderInputSize = 10
req.EncoderMatchedSize = 4

withoutEncoder := &MetricsResponse{Coefficients: &ModelCoefficients{
TTFTIntercept: 100,
TTFTCoeffs: map[string]float64{},
TPOTCoeffs: map[string]float64{},
}}
resp, err := predictor.predictBayesianRidge(req, withoutEncoder, 0.9, ObjectiveMean)
require.NoError(t, err)
assert.Equal(t, 100.0, resp.TTFT)

withEncoder := &MetricsResponse{Coefficients: &ModelCoefficients{
TTFTIntercept: 100,
TTFTCoeffs: map[string]float64{
"encoder_input_size": 2,
"encoder_matched_size": -1,
},
TPOTCoeffs: map[string]float64{},
}}
resp, err = predictor.predictBayesianRidge(req, withEncoder, 0.9, ObjectiveMean)
require.NoError(t, err)
assert.Equal(t, 100.0+2*10-1*4, resp.TTFT)
}
Original file line number Diff line number Diff line change
Expand Up @@ -221,6 +221,12 @@ func (p *Predictor) ValidateTrainingEntry(entry TrainingEntry) error {
if entry.PrefixCacheScore < 0.0 || entry.PrefixCacheScore > 1.0 {
return fmt.Errorf("prefix_cache_score must be between 0.0 and 1.0, got %f", entry.PrefixCacheScore)
}
if entry.EncoderInputSize < 0 {
return fmt.Errorf("encoder_input_size must be non-negative, got %d", entry.EncoderInputSize)
}
if entry.EncoderMatchedSize < 0 || entry.EncoderMatchedSize > entry.EncoderInputSize {
return fmt.Errorf("encoder_matched_size must be between 0 and encoder_input_size (%d), got %d", entry.EncoderInputSize, entry.EncoderMatchedSize)
}
return nil
}

Expand Down
Loading
Loading