Skip to content

Commit f2fe1de

Browse files
jia-gaojia-gao
authored
Use unweighted cached-block count for prefix PD decider (llm-d#1413)
* Use unweighted cached-block count for prefix PD decider The precise prefix cache scorer stores a device-tier-weighted longest-prefix score in PrefixCacheMatchInfo.matchBlocks (RAM tier weighted 0.8 by default). The prefix-based PD decider multiplied that score by the block size and treated the result as a literal cached-token count, so a mostly-RAM-cached prompt had its cached length undercounted and was misrouted to remote prefill — the opposite of what a CPU-offload (LMCache) deployment wants. A short RAM-only hit rounds to zero blocks under int() truncation, dropping the signal entirely. Add an unweighted cachedBlockCount to PrefixCacheMatchInfo (defaulting to matchBlocks so existing producers and consumers are unchanged) and populate it at the two precise-prefix-cache producer sites by counting the contiguous prefix blocks held on each endpoint, independent of device tier. The decider now derives cached tokens from cachedBlockCount. matchBlocks keeps its tier-weighted value for the ranking/affinity consumers, so endpoint ranking is unchanged. Refs: llm-d#1047 Signed-off-by: Jiazhou Gao <gjz140103@gmail.com> * Simplify matchedBlockCount with slices.ContainsFunc Use slices.ContainsFunc for the per-key pod-presence check instead of a manual inner loop, and keep the WithCachedBlockCount chain on one line at the call sites. Review feedback; no behavior change. Signed-off-by: Jiazhou Gao <gjz140103@gmail.com> --------- Signed-off-by: Jiazhou Gao <gjz140103@gmail.com>
1 parent 604da5a commit f2fe1de

7 files changed

Lines changed: 255 additions & 10 deletions

File tree

‎pkg/epp/framework/plugins/datalayer/attribute/prefix/data_types.go‎

Lines changed: 30 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -25,22 +25,38 @@ import (
2525
var PrefixCacheMatchInfoDataKey = plugin.NewDataKey("PrefixCacheMatchInfoDataKey", approxprefixconstants.ApproxPrefixCachePluginType)
2626

2727
type PrefixCacheMatchInfo struct {
28-
// matched prefix length in blocks
28+
// matched prefix length in blocks. For the precise prefix cache this is the
29+
// device-tier-weighted longest-prefix score (e.g. RAM-tier blocks count as
30+
// less than 1.0), suitable for relative endpoint ranking.
2931
matchBlocks int
3032
// total length in blocks
3133
totalBlocks int
3234
// block length in tokens
3335
blockSizeTokens int
36+
// unweighted count of contiguous cached prefix blocks on the endpoint.
37+
// Unlike matchBlocks this is the literal number of cached blocks regardless
38+
// of device tier, so consumers that convert blocks to a token count (e.g.
39+
// the prefix-based PD decider) get an accurate cached-token figure rather
40+
// than a tier-attenuated one. Defaults to matchBlocks when not set.
41+
cachedBlockCount int
3442
}
3543

3644
func NewPrefixCacheMatchInfo(matchBlocks int, totalBlocks int, blockSizeTokens int) *PrefixCacheMatchInfo {
3745
return &PrefixCacheMatchInfo{
38-
matchBlocks: matchBlocks,
39-
totalBlocks: totalBlocks,
40-
blockSizeTokens: blockSizeTokens,
46+
matchBlocks: matchBlocks,
47+
totalBlocks: totalBlocks,
48+
blockSizeTokens: blockSizeTokens,
49+
cachedBlockCount: matchBlocks,
4150
}
4251
}
4352

53+
// WithCachedBlockCount sets the unweighted contiguous cached-block count and
54+
// returns the receiver for chaining.
55+
func (p *PrefixCacheMatchInfo) WithCachedBlockCount(cachedBlockCount int) *PrefixCacheMatchInfo {
56+
p.cachedBlockCount = cachedBlockCount
57+
return p
58+
}
59+
4460
func (p *PrefixCacheMatchInfo) MatchBlocks() int {
4561
return p.matchBlocks
4662
}
@@ -53,10 +69,17 @@ func (p *PrefixCacheMatchInfo) BlockSizeTokens() int {
5369
return p.blockSizeTokens
5470
}
5571

72+
// CachedBlockCount returns the unweighted count of contiguous cached prefix
73+
// blocks on the endpoint.
74+
func (p *PrefixCacheMatchInfo) CachedBlockCount() int {
75+
return p.cachedBlockCount
76+
}
77+
5678
func (p *PrefixCacheMatchInfo) Clone() fwkdl.Cloneable {
5779
return &PrefixCacheMatchInfo{
58-
matchBlocks: p.matchBlocks,
59-
totalBlocks: p.totalBlocks,
60-
blockSizeTokens: p.blockSizeTokens,
80+
matchBlocks: p.matchBlocks,
81+
totalBlocks: p.totalBlocks,
82+
blockSizeTokens: p.blockSizeTokens,
83+
cachedBlockCount: p.cachedBlockCount,
6184
}
6285
}

‎pkg/epp/framework/plugins/datalayer/attribute/prefix/data_types_test.go‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
/*
2+
Copyright 2025 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 prefix
18+
19+
import (
20+
"testing"
21+
22+
"github.com/stretchr/testify/assert"
23+
"github.com/stretchr/testify/require"
24+
)
25+
26+
func TestPrefixCacheMatchInfo_CachedBlockCountDefaultsToMatchBlocks(t *testing.T) {
27+
info := NewPrefixCacheMatchInfo(5, 10, 16)
28+
assert.Equal(t, 5, info.MatchBlocks())
29+
assert.Equal(t, 10, info.TotalBlocks())
30+
assert.Equal(t, 16, info.BlockSizeTokens())
31+
// Unset cachedBlockCount mirrors matchBlocks so existing producers and
32+
// consumers keep their current behavior.
33+
assert.Equal(t, 5, info.CachedBlockCount())
34+
}
35+
36+
func TestPrefixCacheMatchInfo_WithCachedBlockCount(t *testing.T) {
37+
info := NewPrefixCacheMatchInfo(192, 256, 16).WithCachedBlockCount(240)
38+
// matchBlocks keeps the tier-weighted ranking value; cachedBlockCount holds
39+
// the unweighted literal count.
40+
assert.Equal(t, 192, info.MatchBlocks())
41+
assert.Equal(t, 240, info.CachedBlockCount())
42+
}
43+
44+
func TestPrefixCacheMatchInfo_CloneCopiesCachedBlockCount(t *testing.T) {
45+
orig := NewPrefixCacheMatchInfo(192, 256, 16).WithCachedBlockCount(240)
46+
clone, ok := orig.Clone().(*PrefixCacheMatchInfo)
47+
require.True(t, ok)
48+
assert.Equal(t, orig.MatchBlocks(), clone.MatchBlocks())
49+
assert.Equal(t, orig.TotalBlocks(), clone.TotalBlocks())
50+
assert.Equal(t, orig.BlockSizeTokens(), clone.BlockSizeTokens())
51+
assert.Equal(t, orig.CachedBlockCount(), clone.CachedBlockCount())
52+
53+
// Mutating the clone must not affect the original.
54+
clone.WithCachedBlockCount(1)
55+
assert.Equal(t, 240, orig.CachedBlockCount())
56+
assert.Equal(t, 1, clone.CachedBlockCount())
57+
}

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -276,8 +276,9 @@ func (p *Producer) Produce(ctx context.Context,
276276
if matchLen > maxMatch {
277277
maxMatch = matchLen
278278
}
279+
cachedBlocks := matchedBlockCount(blockKeys, keyToPods, addr)
279280
ep.Put(p.dk.String(),
280-
attrprefix.NewPrefixCacheMatchInfo(matchLen, totalBlocks, p.blockSizeTokens))
281+
attrprefix.NewPrefixCacheMatchInfo(matchLen, totalBlocks, p.blockSizeTokens).WithCachedBlockCount(cachedBlocks))
281282
}
282283

283284
if p.speculativeEnabled {

‎pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/utils.go‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,9 @@ package preciseprefixcache
1818

1919
import (
2020
"fmt"
21+
"slices"
2122

23+
"github.com/llm-d/llm-d-kv-cache/pkg/kvcache/kvblock"
2224
"k8s.io/apimachinery/pkg/util/sets"
2325

2426
"github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling"
@@ -36,3 +38,19 @@ func extractEndpointSet(endpoints []scheduling.Endpoint) sets.Set[string] {
3638
}
3739
return endpointSet
3840
}
41+
42+
// matchedBlockCount returns the number of contiguous cached prefix blocks held
43+
// by podID, counting from the first block until the first block the pod does
44+
// not hold. This is the unweighted counterpart of the device-tier-weighted
45+
// kvblock scorer: every cached block counts as one regardless of device tier,
46+
// so a pod present at keys[0..n-1] yields n.
47+
func matchedBlockCount(keys []kvblock.BlockHash, keyToPods map[kvblock.BlockHash][]kvblock.PodEntry, podID string) int {
48+
count := 0
49+
for _, key := range keys {
50+
if !slices.ContainsFunc(keyToPods[key], func(e kvblock.PodEntry) bool { return e.PodIdentifier == podID }) {
51+
break
52+
}
53+
count++
54+
}
55+
return count
56+
}

‎pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/utils_test.go‎

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
/*
2+
Copyright 2026 The llm-d 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 preciseprefixcache
18+
19+
import (
20+
"testing"
21+
22+
"github.com/llm-d/llm-d-kv-cache/pkg/kvcache/kvblock"
23+
"github.com/stretchr/testify/assert"
24+
)
25+
26+
func TestMatchedBlockCount(t *testing.T) {
27+
const (
28+
podA = "10.0.0.1:8000"
29+
podB = "10.0.0.2:8000"
30+
)
31+
keys := []kvblock.BlockHash{1, 2, 3, 4}
32+
33+
// gpu/cpu tiers must count identically — the unweighted count ignores tier.
34+
gpu := func(pod string) kvblock.PodEntry { return kvblock.PodEntry{PodIdentifier: pod, DeviceTier: "gpu"} }
35+
cpu := func(pod string) kvblock.PodEntry { return kvblock.PodEntry{PodIdentifier: pod, DeviceTier: "cpu"} }
36+
37+
tests := []struct {
38+
name string
39+
keyToPods map[kvblock.BlockHash][]kvblock.PodEntry
40+
podID string
41+
want int
42+
}{
43+
{
44+
name: "all blocks held on RAM/cpu tier count fully (unweighted)",
45+
keyToPods: map[kvblock.BlockHash][]kvblock.PodEntry{
46+
1: {cpu(podA)}, 2: {cpu(podA)}, 3: {cpu(podA)}, 4: {cpu(podA)},
47+
},
48+
podID: podA,
49+
want: 4,
50+
},
51+
{
52+
name: "single RAM block counts as one, not zero",
53+
keyToPods: map[kvblock.BlockHash][]kvblock.PodEntry{
54+
1: {cpu(podA)},
55+
},
56+
podID: podA,
57+
want: 1,
58+
},
59+
{
60+
name: "stops at first missing block",
61+
keyToPods: map[kvblock.BlockHash][]kvblock.PodEntry{
62+
1: {gpu(podA)}, 2: {gpu(podA)}, 4: {gpu(podA)}, // block 3 missing
63+
},
64+
podID: podA,
65+
want: 2,
66+
},
67+
{
68+
name: "pod absent from first block yields zero",
69+
keyToPods: map[kvblock.BlockHash][]kvblock.PodEntry{
70+
1: {gpu(podB)}, 2: {gpu(podA)},
71+
},
72+
podID: podA,
73+
want: 0,
74+
},
75+
{
76+
name: "counts are per-pod independent",
77+
keyToPods: map[kvblock.BlockHash][]kvblock.PodEntry{
78+
1: {gpu(podA), cpu(podB)}, 2: {gpu(podA)}, 3: {cpu(podB)},
79+
},
80+
podID: podA,
81+
want: 2,
82+
},
83+
{
84+
name: "empty index yields zero",
85+
keyToPods: map[kvblock.BlockHash][]kvblock.PodEntry{},
86+
podID: podA,
87+
want: 0,
88+
},
89+
}
90+
91+
for _, tt := range tests {
92+
t.Run(tt.name, func(t *testing.T) {
93+
assert.Equal(t, tt.want, matchedBlockCount(keys, tt.keyToPods, tt.podID))
94+
})
95+
}
96+
}

‎pkg/epp/framework/plugins/scheduling/profilehandler/disagg/prefix_based_pd_decider.go‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -127,8 +127,11 @@ func (d *PrefixBasedPDDecider) disaggregate(ctx context.Context, request *schedu
127127
return false
128128
}
129129

130-
// number of cached tokens
131-
hitPrefixTokens := prefixCacheMatchInfo.MatchBlocks() * prefixCacheMatchInfo.BlockSizeTokens()
130+
// number of cached tokens. Use the unweighted cached-block count, not the
131+
// tier-weighted match score: a RAM-cached prefix must contribute its full
132+
// token count here, otherwise the non-cached suffix is overestimated and
133+
// requests with large local-RAM hits are misrouted to remote prefill.
134+
hitPrefixTokens := prefixCacheMatchInfo.CachedBlockCount() * prefixCacheMatchInfo.BlockSizeTokens()
132135
// length of non-cached suffix in tokens
133136
nonCachedTokens := inputTokens - hitPrefixTokens
134137

‎pkg/epp/framework/plugins/scheduling/profilehandler/disagg/prefix_based_pd_decider_test.go‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -354,6 +354,53 @@ func TestDisaggregate(t *testing.T) {
354354
}
355355
}
356356

357+
// TestDisaggregate_UsesUnweightedCachedBlockCount reproduces the #1047
358+
// RAM-cache misrouting scenario. The precise prefix cache scorer stores a
359+
// device-tier-weighted match score in matchBlocks (RAM tier = 0.8), but the
360+
// literal cached-block count lives in cachedBlockCount. The decider must use
361+
// the unweighted count so a mostly-RAM-cached prompt is not pushed onto the
362+
// remote-prefill path.
363+
//
364+
// Issue parameters: blockSize=16, inputTokens=4096, a 240-block contiguous hit
365+
// (3840 cached tokens, real non-cached suffix 256), threshold 512.
366+
func TestDisaggregate_UsesUnweightedCachedBlockCount(t *testing.T) {
367+
ctx := utils.NewTestContext(t)
368+
369+
const (
370+
blockSize = 16
371+
inputTokens = 4096
372+
totalBlocks = inputTokens / blockSize // 256
373+
cachedBlocks = 240 // contiguous hit (3840 tokens)
374+
ramWeightedScore = 192 // int(240 * 0.8) as stored in matchBlocks
375+
nonCachedTokens = 512
376+
)
377+
378+
newEndpoint := func(info *attrprefix.PrefixCacheMatchInfo) scheduling.Endpoint {
379+
ep := makeTestEndpointBase()
380+
ep.Put(attrprefix.PrefixCacheMatchInfoDataKey.String(), info)
381+
return ep
382+
}
383+
// Exact token count via the tokenized-prompt path the decider reads.
384+
req := withTokens(completionsRequestWithPrompt(fwkrh.Prompt{}), inputTokens)
385+
386+
decider, err := NewPrefixBasedPDDecider(PrefixBasedPDDeciderConfig{NonCachedTokens: nonCachedTokens})
387+
require.NoError(t, err)
388+
389+
// Fixed behavior: cachedBlockCount carries the true 240 blocks, so
390+
// nonCached = 4096 - 240*16 = 256 < 512 → decode-only (no remote prefill).
391+
fixed := newEndpoint(attrprefix.NewPrefixCacheMatchInfo(ramWeightedScore, totalBlocks, blockSize).
392+
WithCachedBlockCount(cachedBlocks))
393+
assert.False(t, decider.disaggregate(ctx, req, fixed),
394+
"RAM-cached prefix must stay decode-only when the unweighted cached-block count is used")
395+
396+
// Buggy behavior guard: if only the tier-weighted score (192) were
397+
// available as the block count, nonCached = 4096 - 192*16 = 1024 >= 512
398+
// would misroute to remote prefill.
399+
weightedOnly := newEndpoint(attrprefix.NewPrefixCacheMatchInfo(ramWeightedScore, totalBlocks, blockSize))
400+
assert.True(t, decider.disaggregate(ctx, req, weightedOnly),
401+
"sanity: the tier-weighted score alone undercounts cached blocks and misroutes")
402+
}
403+
357404
func TestDisaggregateNoPrefixInfo(t *testing.T) {
358405
ctx := utils.NewTestContext(t)
359406

0 commit comments

Comments
 (0)