Skip to content

Commit d002e73

Browse files
committed
feat(prefix): auto-tune prefix match cap from model max_model_len
The approximate-prefix-cache producer caps prefix matching at a fixed maxPrefixTokensToMatch (default 128k). On models with a larger context window that under-matches long shared prefixes, and it's one more knob to set per model. When autoTune is on and maxPrefixTokensToMatch is left at its default, derive the cap from each endpoint's max_model_len instead. The value is read from the model server's /v1/models, fetched once per pod off the request path and cached. Until it resolves (or if the endpoint can't be reached) we fall back to the existing default, so behaviour is unchanged in the common case. An explicitly configured cap is always respected. The fetch reuses the model-server-metrics scheme/TLS flags so /v1/models follows the same protocol as /metrics. Signed-off-by: Mayur Das <mayur.das@neevcloud.com>
1 parent 546ea56 commit d002e73

6 files changed

Lines changed: 279 additions & 2 deletions

File tree

pkg/epp/framework/plugins/datalayer/attribute/models/data_types.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,9 @@ type ModelDataCollection []ModelData
3737
type ModelData struct {
3838
ID string `json:"id"`
3939
Parent string `json:"parent,omitempty"`
40+
// MaxModelLen is the model's context window in tokens as reported by the
41+
// model server; zero when the server does not report it.
42+
MaxModelLen int `json:"max_model_len,omitempty"`
4043
}
4144

4245
// String returns a string representation of the model info
Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
/*
2+
Copyright 2026 The Kubernetes Authors.
3+
4+
Licensed under the Apache License, Version 2.0 (the "License");
5+
you may not use this file except in compliance with the License.
6+
You may obtain a copy of the License at
7+
8+
http://www.apache.org/licenses/LICENSE-2.0
9+
10+
Unless required by applicable law or agreed to in writing, software
11+
distributed under the License is distributed on an "AS IS" BASIS,
12+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
See the License for the specific language governing permissions and
14+
limitations under the License.
15+
*/
16+
17+
package models
18+
19+
import (
20+
"testing"
21+
22+
"github.com/stretchr/testify/assert"
23+
)
24+
25+
// TestModelDataCollectionClone verifies Clone copies entries (including
26+
// MaxModelLen) and is independent of the original.
27+
func TestModelDataCollectionClone(t *testing.T) {
28+
assert.Nil(t, ModelDataCollection(nil).Clone())
29+
30+
orig := ModelDataCollection{{ID: "m1", MaxModelLen: 100}, {ID: "m2", MaxModelLen: 200}}
31+
cloned, ok := orig.Clone().(ModelDataCollection)
32+
assert.True(t, ok)
33+
assert.Equal(t, orig, cloned)
34+
35+
// Mutating the clone must not affect the original.
36+
cloned[0].MaxModelLen = 999
37+
assert.Equal(t, 100, orig[0].MaxModelLen)
38+
}

pkg/epp/framework/plugins/datalayer/extractor/models/extractor_test.go

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,3 +106,32 @@ func TestExtractorExtract(t *testing.T) {
106106
})
107107
}
108108
}
109+
110+
// TestExtractorCapturesMaxModelLen verifies the extractor surfaces max_model_len
111+
// from the /v1/models payload onto the endpoint attribute.
112+
func TestExtractorCapturesMaxModelLen(t *testing.T) {
113+
ctx := context.Background()
114+
extPlugin, err := ModelServerExtractorFactory("test-extractor", nil, nil)
115+
if err != nil {
116+
t.Fatalf("failed to create extractor: %v", err)
117+
}
118+
extractor := extPlugin.(*ModelExtractor)
119+
ep := fwkdl.NewEndpoint(nil, nil)
120+
121+
resp := &ModelResponse{Data: []attrmodels.ModelData{{ID: "m", MaxModelLen: 131072}}}
122+
if err := extractor.Extract(ctx, fwkdl.PollInput[*ModelResponse]{Payload: resp, Endpoint: ep}); err != nil {
123+
t.Fatalf("unexpected error: %v", err)
124+
}
125+
126+
val, ok := ep.GetAttributes().Get(attrmodels.ModelsAttributeKey.String())
127+
if !ok {
128+
t.Fatal("expected models attribute to be set")
129+
}
130+
models, ok := val.(attrmodels.ModelDataCollection)
131+
if !ok {
132+
t.Fatalf("unexpected attribute type %T", val)
133+
}
134+
if len(models) != 1 || models[0].MaxModelLen != 131072 {
135+
t.Errorf("expected max_model_len 131072, got %+v", models)
136+
}
137+
}
Lines changed: 167 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,167 @@
1+
/*
2+
Copyright 2026 The Kubernetes Authors.
3+
4+
Licensed under the Apache License, Version 2.0 (the "License");
5+
you may not use this file except in compliance with the License.
6+
You may obtain a copy of the License at
7+
8+
http://www.apache.org/licenses/LICENSE-2.0
9+
10+
Unless required by applicable law or agreed to in writing, software
11+
distributed under the License is distributed on an "AS IS" BASIS,
12+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
See the License for the specific language governing permissions and
14+
limitations under the License.
15+
*/
16+
17+
package approximateprefix
18+
19+
import (
20+
"context"
21+
"testing"
22+
23+
"github.com/google/uuid"
24+
"github.com/stretchr/testify/assert"
25+
k8stypes "k8s.io/apimachinery/pkg/types"
26+
27+
fwkdl "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/datalayer"
28+
"github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin"
29+
fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling"
30+
attrmodels "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/models"
31+
attrprefix "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/prefix"
32+
)
33+
34+
// modelLenEndpoint builds an endpoint whose /v1/models attribute reports the
35+
// given max_model_len values, one per model entry. With no values, the attribute
36+
// is left unset.
37+
func modelLenEndpoint(maxModelLens ...int) fwksched.Endpoint {
38+
attrs := fwkdl.NewAttributes()
39+
if maxModelLens != nil {
40+
models := make(attrmodels.ModelDataCollection, len(maxModelLens))
41+
for i, ml := range maxModelLens {
42+
models[i] = attrmodels.ModelData{ID: "m", MaxModelLen: ml}
43+
}
44+
attrs.Put(attrmodels.ModelsAttributeKey.String(), models)
45+
}
46+
return fwksched.NewEndpoint(&fwkdl.EndpointMetadata{NamespacedName: k8stypes.NamespacedName{Name: "pod1"}}, fwkdl.NewMetrics(), attrs)
47+
}
48+
49+
// TestProduceAutoTunesFromModelLen verifies the prefix cap is derived from the
50+
// endpoint's max_model_len attribute when auto-tuning, and falls back otherwise.
51+
func TestProduceAutoTunesFromModelLen(t *testing.T) {
52+
disableMinBlockSizeClamp(t)
53+
54+
// A 20-block prompt at block size 1.
55+
tokens := make([]uint32, 20)
56+
for i := range tokens {
57+
tokens[i] = uint32(i + 1)
58+
}
59+
60+
produceBlocks := func(t *testing.T, p *dataProducer, ep fwksched.Endpoint) int {
61+
t.Helper()
62+
req := &fwksched.InferenceRequest{RequestID: uuid.NewString(), TargetModel: "m", Body: tokenizedBody(tokens)}
63+
assert.NoError(t, p.Produce(context.Background(), req, []fwksched.Endpoint{ep}))
64+
state, err := plugin.ReadPluginStateKey[*SchedulingContextState](
65+
p.PluginState(), req.RequestID, plugin.StateKey(ApproxPrefixCachePluginType))
66+
assert.NoError(t, err)
67+
return len(state.PerPromptHashes[0])
68+
}
69+
70+
autoTuneConfig := config{
71+
AutoTune: true,
72+
BlockSizeTokens: 1,
73+
MaxPrefixBlocksToMatch: defaultMaxPrefixBlocks,
74+
MaxPrefixTokensToMatch: defaultMaxPrefixTokens, // left at default -> auto-tune on
75+
LRUCapacityPerServer: defaultLRUCapacityPerServer,
76+
}
77+
78+
t.Run("caps at max_model_len when known", func(t *testing.T) {
79+
p, err := newDataProducer(context.Background(), ApproxPrefixCachePluginType, autoTuneConfig, testHandle())
80+
assert.NoError(t, err)
81+
assert.True(t, p.config.autoTuneMaxPrefixTokens)
82+
// max_model_len = 10 -> cap = 10/blockSize(1) = 10 blocks.
83+
assert.Equal(t, 10, produceBlocks(t, p, modelLenEndpoint(10)))
84+
})
85+
86+
t.Run("uses largest across models", func(t *testing.T) {
87+
p, err := newDataProducer(context.Background(), ApproxPrefixCachePluginType, autoTuneConfig, testHandle())
88+
assert.NoError(t, err)
89+
// Largest wins regardless of position (first entry here).
90+
assert.Equal(t, 15, produceBlocks(t, p, modelLenEndpoint(15, 8, 12)))
91+
})
92+
93+
t.Run("falls back on negative max_model_len", func(t *testing.T) {
94+
p, err := newDataProducer(context.Background(), ApproxPrefixCachePluginType, autoTuneConfig, testHandle())
95+
assert.NoError(t, err)
96+
assert.Equal(t, 20, produceBlocks(t, p, modelLenEndpoint(-5)))
97+
})
98+
99+
t.Run("falls back on empty models collection", func(t *testing.T) {
100+
p, err := newDataProducer(context.Background(), ApproxPrefixCachePluginType, autoTuneConfig, testHandle())
101+
assert.NoError(t, err)
102+
// Attribute present but no model entries.
103+
attrs := fwkdl.NewAttributes()
104+
attrs.Put(attrmodels.ModelsAttributeKey.String(), attrmodels.ModelDataCollection{})
105+
ep := fwksched.NewEndpoint(&fwkdl.EndpointMetadata{NamespacedName: k8stypes.NamespacedName{Name: "pod1"}}, fwkdl.NewMetrics(), attrs)
106+
assert.Equal(t, 20, produceBlocks(t, p, ep))
107+
})
108+
109+
t.Run("falls back to default when attribute absent", func(t *testing.T) {
110+
p, err := newDataProducer(context.Background(), ApproxPrefixCachePluginType, autoTuneConfig, testHandle())
111+
assert.NoError(t, err)
112+
// No models attribute -> default cap (131072) far exceeds 20 blocks.
113+
assert.Equal(t, 20, produceBlocks(t, p, modelLenEndpoint()))
114+
})
115+
116+
t.Run("falls back when max_model_len is zero", func(t *testing.T) {
117+
p, err := newDataProducer(context.Background(), ApproxPrefixCachePluginType, autoTuneConfig, testHandle())
118+
assert.NoError(t, err)
119+
// Attribute present but server did not report the field.
120+
assert.Equal(t, 20, produceBlocks(t, p, modelLenEndpoint(0)))
121+
})
122+
123+
t.Run("explicit cap is not auto-tuned", func(t *testing.T) {
124+
explicit := autoTuneConfig
125+
explicit.MaxPrefixTokensToMatch = 5 // operator-pinned, non-default
126+
p, err := newDataProducer(context.Background(), ApproxPrefixCachePluginType, explicit, testHandle())
127+
assert.NoError(t, err)
128+
assert.False(t, p.config.autoTuneMaxPrefixTokens)
129+
// Even with a known model length, the explicit cap (5) wins.
130+
assert.Equal(t, 5, produceBlocks(t, p, modelLenEndpoint(10)))
131+
})
132+
133+
t.Run("empty pods uses default cap", func(t *testing.T) {
134+
p, err := newDataProducer(context.Background(), ApproxPrefixCachePluginType, autoTuneConfig, testHandle())
135+
assert.NoError(t, err)
136+
137+
req := &fwksched.InferenceRequest{RequestID: uuid.NewString(), TargetModel: "m", Body: tokenizedBody(tokens)}
138+
assert.NoError(t, p.Produce(context.Background(), req, []fwksched.Endpoint{}))
139+
state, err := plugin.ReadPluginStateKey[*SchedulingContextState](
140+
p.PluginState(), req.RequestID, plugin.StateKey(ApproxPrefixCachePluginType))
141+
assert.NoError(t, err)
142+
assert.Equal(t, 20, len(state.PerPromptHashes[0]))
143+
})
144+
}
145+
146+
// TestConsumesDeclaresModelsDependency verifies the producer declares the
147+
// /v1/models attribute as an optional dependency so the cap can be auto-tuned
148+
// when present and fall back when absent.
149+
func TestConsumesDeclaresModelsDependency(t *testing.T) {
150+
p, err := newDataProducer(context.Background(), ApproxPrefixCachePluginType, defaultConfig, testHandle())
151+
assert.NoError(t, err)
152+
_, ok := p.Consumes().Optional[attrmodels.ModelsAttributeKey]
153+
assert.True(t, ok, "models attribute must be an optional dependency")
154+
}
155+
156+
// TestMaxModelLenHelper covers the attribute reader's non-collection and
157+
// missing-attribute paths.
158+
func TestMaxModelLenHelper(t *testing.T) {
159+
// Missing attribute -> 0.
160+
assert.Equal(t, 0, maxModelLen(modelLenEndpoint()))
161+
162+
// Wrong type stored under the key -> 0.
163+
attrs := fwkdl.NewAttributes()
164+
attrs.Put(attrmodels.ModelsAttributeKey.String(), attrprefix.NewPrefixCacheMatchInfo(0, 0, 0))
165+
ep := fwksched.NewEndpoint(&fwkdl.EndpointMetadata{NamespacedName: k8stypes.NamespacedName{Name: "pod1"}}, fwkdl.NewMetrics(), attrs)
166+
assert.Equal(t, 0, maxModelLen(ep))
167+
}

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

Lines changed: 37 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ import (
3030
"github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin"
3131
"github.com/llm-d/llm-d-router/pkg/epp/framework/interface/requestcontrol"
3232
fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling"
33+
attrmodels "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/models"
3334
attrprefix "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/prefix"
3435
approxprefixconstants "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/constants"
3536
tokenproducer "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer"
@@ -80,6 +81,9 @@ func (p *dataProducer) Produces() map[plugin.DataKey]any {
8081
func (p *dataProducer) Consumes() plugin.DataDependencies {
8182
return plugin.DataDependencies{
8283
Required: map[plugin.DataKey]any{tokenproducer.TokenizedPromptDataKey: fwksched.TokenizedPrompt{}},
84+
// Optional: supplies each endpoint's max_model_len for auto-tuning the cap.
85+
// Produce falls back to the default cap when it is absent.
86+
Optional: map[plugin.DataKey]any{attrmodels.ModelsAttributeKey: attrmodels.ModelDataCollection{}},
8387
}
8488
}
8589

@@ -118,6 +122,10 @@ func newDataProducer(ctx context.Context, name string, config config, handle plu
118122

119123
indexer := newIndexer(ctx, config.LRUCapacityPerServer, name, ApproxPrefixCachePluginType)
120124

125+
// Auto-tune the cap from max_model_len only when AutoTune is on and the
126+
// operator left MaxPrefixTokensToMatch at its default (no explicit cap).
127+
config.autoTuneMaxPrefixTokens = config.AutoTune && config.MaxPrefixTokensToMatch == defaultMaxPrefixTokens
128+
121129
p := &dataProducer{
122130
typedName: plugin.TypedName{
123131
Type: ApproxPrefixCachePluginType,
@@ -175,9 +183,16 @@ func (p *dataProducer) PluginState() *plugin.PluginState {
175183
// Produce is called by the director before scheduling requests.
176184
func (p *dataProducer) Produce(ctx context.Context, request *fwksched.InferenceRequest, pods []fwksched.Endpoint) error {
177185
blockSize := p.GetBlockSize(pods)
186+
maxTokens := p.config.MaxPrefixTokensToMatch
187+
// Raise the cap to the model's full context window when max_model_len is known.
188+
if p.config.autoTuneMaxPrefixTokens && len(pods) > 0 {
189+
if modelLen := maxModelLen(pods[0]); modelLen > 0 {
190+
maxTokens = modelLen
191+
}
192+
}
178193
maxBlocks := p.config.MaxPrefixBlocksToMatch
179-
if p.config.MaxPrefixTokensToMatch > 0 && blockSize > 0 {
180-
maxBlocks = p.config.MaxPrefixTokensToMatch / blockSize
194+
if maxTokens > 0 && blockSize > 0 {
195+
maxBlocks = maxTokens / blockSize
181196
}
182197
perPromptHashes := getBlockHashes(ctx, request, blockSize, maxBlocks)
183198

@@ -307,6 +322,26 @@ func (p *dataProducer) GetBlockSize(endpoints []fwksched.Endpoint) int {
307322
return blockSize
308323
}
309324

325+
// maxModelLen returns the largest max_model_len reported for ep via the
326+
// /v1/models attribute, or 0 when absent or non-positive.
327+
func maxModelLen(ep fwksched.Endpoint) int {
328+
val, ok := ep.Get(attrmodels.ModelsAttributeKey.String())
329+
if !ok {
330+
return 0
331+
}
332+
models, ok := val.(attrmodels.ModelDataCollection)
333+
if !ok {
334+
return 0
335+
}
336+
maxLen := 0
337+
for i := range models {
338+
if models[i].MaxModelLen > maxLen {
339+
maxLen = models[i].MaxModelLen
340+
}
341+
}
342+
return maxLen
343+
}
344+
310345
// ApproxPrefixCacheFactory is the factory function for the prefix cache data producer plugin.
311346
func ApproxPrefixCacheFactory(name string, rawParameters *json.Decoder, handle plugin.Handle) (plugin.Plugin, error) {
312347
parameters := defaultConfig

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -140,6 +140,11 @@ type config struct {
140140
MaxPrefixTokensToMatch int `json:"maxPrefixTokensToMatch"`
141141
// Max capacity size of the LRU indexer in number of entries per server (pod).
142142
LRUCapacityPerServer int `json:"lruCapacityPerServer"`
143+
144+
// autoTuneMaxPrefixTokens, set internally (never from JSON), derives the cap
145+
// per endpoint from the model's max_model_len. Enabled when AutoTune is on and
146+
// MaxPrefixTokensToMatch was left at its default (no operator-pinned cap).
147+
autoTuneMaxPrefixTokens bool
143148
}
144149

145150
// defaultConfig provides sensible defaults for the prefix cache plugins.

0 commit comments

Comments
 (0)