fix(epp): widen p2p source sampling to a one-block band

Exact-count argmax re-concentrates pull traffic the moment one replica
drifts a block ahead: the queue weighting never sees near-tied peers,
so a deeply-queued source absorbs every pull while an idle peer one
block short sits unused. Sample among peers within one block of the
maximum instead; the one-block shortfall costs the destination a
single block of recompute, noise next to a queue-depth difference on
the source. The stashed count is the chosen peer's own, so the
minCachedTokenDelta gate judges the actual pull benefit.

Signed-off-by: nilig <nili.ifergan@gmail.com>

Author: nilig

pkg/epp/framework/plugins/requestcontrol/dataproducer/p2psource/README.md

@@ -2,9 +2,9 @@
 
 **Type:** `p2p-source-producer`
 
-Sets the `x-kv-cache-source-host-port` header to the endpoint holding the most cached prompt prefix, so the routing sidecar can pull those blocks over the P2P connector instead of recomputing them. Runs in the request handling's `DataProducer` phase before scheduling, then emits the header in `PreRequest` after the scheduling decision.
+Sets the `x-kv-cache-source-host-port` header to an endpoint within one block of the most cached prompt prefix, so the routing sidecar can pull those blocks over the P2P connector instead of recomputing them. Runs in the request handling's `DataProducer` phase before scheduling, then emits the header in `PreRequest` after the scheduling decision.
 
-For each request the plugin consumes the per-endpoint `PrefixCacheMatchInfo` produced by a prefix-cache producer (`approx-prefix-cache-producer` or `precise-prefix-cache-producer`) and selects the source among the endpoints caching the most prompt tokens, sampling one with probability proportional to `1/(1+waiting queue)` using a request-ID hash as the deterministic sampling coordinate. Load-blind argmax alone would send every consumer of a widely-replicated prefix to the same peer, concentrating pull traffic on one source; a hard minimum-queue rule would herd whole scrape windows onto whichever peer last read the shortest queue, since metrics refresh about once per second and a one-request difference is noise at that granularity. Proportional weighting spreads exact ties uniformly, shifts share about 2:1 on a one-request queue difference, and effectively starves deeply-queued sources. The waiting queue is a proxy for engine-step responsiveness (a busy engine serves its P2P session slower), not for pull-serve fanout; the hash spread is what distributes pull traffic among equally-suitable sources. After scheduling it compares the chosen peer against the pod that will compute the prefix — the `prefill` profile target under P/D disaggregation, otherwise the primary target — and sets the header only when the peer out-caches the computing pod by at least `minCachedTokenDelta` tokens. Any inbound value of the header is removed. When no peer out-caches the computing pod, the request proceeds unchanged.
+For each request the plugin consumes the per-endpoint `PrefixCacheMatchInfo` produced by a prefix-cache producer (`approx-prefix-cache-producer` or `precise-prefix-cache-producer`) and selects the source among the endpoints caching within one block of the most prompt tokens, sampling one with probability proportional to `1/(1+waiting queue)` using a request-ID hash as the deterministic sampling coordinate. Load-blind argmax alone would send every consumer of a widely-replicated prefix to the same peer, concentrating pull traffic on one source; a hard minimum-queue rule would herd whole scrape windows onto whichever peer last read the shortest queue, since metrics refresh about once per second and a one-request difference is noise at that granularity. The one-block band keeps a peer one block short competitive — the shortfall costs the destination a single block of recompute, noise next to a queue-depth difference on the source — so exact-count argmax cannot re-concentrate traffic the moment one replica drifts a block ahead. Proportional weighting spreads ties uniformly, shifts share about 2:1 on a one-request queue difference, and effectively starves deeply-queued sources. The waiting queue is a proxy for engine-step responsiveness (a busy engine serves its P2P session slower), not for pull-serve fanout; the hash spread is what distributes pull traffic among equally-suitable sources. After scheduling it compares the chosen peer against the pod that will compute the prefix — the `prefill` profile target under P/D disaggregation, otherwise the primary target — and sets the header only when the peer out-caches the computing pod by at least `minCachedTokenDelta` tokens. Any inbound value of the header is removed. When no peer out-caches the computing pod, the request proceeds unchanged.
 
 **Parameters:**
 

pkg/epp/framework/plugins/requestcontrol/dataproducer/p2psource/producer.go

@@ -14,9 +14,10 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-// Package p2psource emits the KV cache source header: the candidate pod
-// holding the most cached prefix KV blocks for the request, for the routing
-// sidecar to pull from over the P2P connector instead of recomputing them.
+// Package p2psource emits the KV cache source header: a candidate pod
+// within one block of the most cached prefix KV blocks for the request, for
+// the routing sidecar to pull from over the P2P connector instead of
+// recomputing them.
 package p2psource
 
 import (
@@ -70,11 +71,11 @@ var (
 	_ requestcontrol.PreRequest   = &Producer{}
 )
 
-// Producer stashes the candidate endpoint holding the most cached prefix
-// tokens during Produce, and in PreRequest sets routing.KVCacheSourceHeader
-// to that peer when it out-caches the pod computing the prefix (the prefill
-// endpoint under P/D disaggregation, the primary endpoint otherwise) by at
-// least minCachedTokenDelta tokens.
+// Producer stashes a pull-source candidate holding within one block of the
+// most cached prefix tokens during Produce, and in PreRequest sets
+// routing.KVCacheSourceHeader to that peer when it out-caches the pod
+// computing the prefix (the prefill endpoint under P/D disaggregation, the
+// primary endpoint otherwise) by at least minCachedTokenDelta tokens.
 type Producer struct {
 	typedName           plugin.TypedName
 	prefixMatchDataKey  plugin.DataKey
@@ -129,9 +130,10 @@ func (p *Producer) Consumes() plugin.DataDependencies {
 	}
 }
 
-// bestMatchPeer is the candidate endpoint holding the most cached prompt
-// tokens for the request, stashed as a request attribute so PreRequest can
-// compare it against the scheduled endpoint.
+// bestMatchPeer is the pull-source candidate chosen during Produce, stashed
+// as a request attribute so PreRequest can compare it against the scheduled
+// endpoint. cachedTokens is the chosen peer's own count, which may be one
+// block below the pool maximum.
 type bestMatchPeer struct {
 	hostPort     string
 	cachedTokens int
@@ -144,17 +146,20 @@ func (p *Producer) attrKey() string {
 }
 
 // Produce reads each candidate's PrefixCacheMatchInfo and stashes the chosen
-// source on the request: among the endpoints holding the most cached prompt
-// tokens, one is sampled with probability proportional to 1/(1+waiting
-// queue), using a request-ID hash as the sampling coordinate. Load-blind
-// argmax alone would send every consumer of a widely-replicated prefix to
-// the same peer (equal counts lose to iteration order), concentrating pull
-// traffic on one source. A hard minimum-queue rule would too: metrics
-// refresh once per scrape interval, so a stale one-request difference would
-// herd a whole window of requests onto the single "idle" peer. Proportional
-// weights spread exact ties uniformly, mildly prefer a one-request-shorter
-// queue, and starve deeply-queued sources. No-op when no candidate holds any
-// cached block.
+// source on the request: among the endpoints within one block of the most
+// cached prompt tokens, one is sampled with probability proportional to
+// 1/(1+waiting queue), using a request-ID hash as the sampling coordinate.
+// Load-blind argmax alone would send every consumer of a widely-replicated
+// prefix to the same peer (equal counts lose to iteration order),
+// concentrating pull traffic on one source. A hard minimum-queue rule would
+// too: metrics refresh once per scrape interval, so a stale one-request
+// difference would herd a whole window of requests onto the single "idle"
+// peer. The one-block band keeps exact-count argmax from re-concentrating
+// the moment one replica drifts a block ahead: a peer one block short costs
+// the destination a single block of recompute, noise next to a queue-depth
+// difference on the source. Proportional weights spread ties uniformly,
+// mildly prefer a one-request-shorter queue, and starve deeply-queued
+// sources. No-op when no candidate holds any cached block.
 func (p *Producer) Produce(ctx context.Context, request *scheduling.InferenceRequest, endpoints []scheduling.Endpoint) error {
 	maxCached := 0
 	for _, ep := range endpoints {
@@ -166,14 +171,17 @@ func (p *Producer) Produce(ctx context.Context, request *scheduling.InferenceReq
 	best := bestMatchPeer{}
 	if maxCached > 0 {
 		var candidates []scheduling.Endpoint
+		var cachedCounts []int
 		var weights []float64
 		total := 0.0
 		for _, ep := range endpoints {
-			if ep.GetMetadata() == nil || p.cachedTokenCount(ep) != maxCached {
+			cached, blockSize := p.cachedTokens(ep)
+			if ep.GetMetadata() == nil || cached == 0 || cached+blockSize < maxCached {
 				continue
 			}
 			w := 1.0 / (1.0 + float64(waitingQueueSize(ep)))
 			candidates = append(candidates, ep)
+			cachedCounts = append(cachedCounts, cached)
 			weights = append(weights, w)
 			total += w
 		}
@@ -190,7 +198,7 @@ func (p *Producer) Produce(ctx context.Context, request *scheduling.InferenceReq
 			md := candidates[chosen].GetMetadata()
 			best = bestMatchPeer{
 				hostPort:     net.JoinHostPort(md.Address, md.Port),
-				cachedTokens: maxCached,
+				cachedTokens: cachedCounts[chosen],
 			}
 		}
 	}
@@ -273,17 +281,24 @@ func (p *Producer) PreRequest(ctx context.Context, request *scheduling.Inference
 	logger.Info("set KV cache source header", "requestID", request.RequestID, "value", best.hostPort)
 }
 
-// cachedTokenCount returns the endpoint's cached prompt tokens (unweighted
-// cached-block count times the block size) from its PrefixCacheMatchInfo,
-// or 0 when absent.
-func (p *Producer) cachedTokenCount(ep scheduling.Endpoint) int {
+// cachedTokens returns the endpoint's cached prompt tokens (unweighted
+// cached-block count times the block size) and its block size from its
+// PrefixCacheMatchInfo, or zeros when absent.
+func (p *Producer) cachedTokens(ep scheduling.Endpoint) (tokens, blockSize int) {
 	raw, ok := ep.Get(p.prefixMatchDataKey.String())
 	if !ok {
-		return 0
+		return 0, 0
 	}
 	info, ok := raw.(*attrprefix.PrefixCacheMatchInfo)
 	if !ok {
-		return 0
+		return 0, 0
 	}
-	return info.CachedBlockCount() * info.BlockSizeTokens()
+	return info.CachedBlockCount() * info.BlockSizeTokens(), info.BlockSizeTokens()
+}
+
+// cachedTokenCount returns the endpoint's cached prompt tokens, or 0 when
+// its PrefixCacheMatchInfo is absent.
+func (p *Producer) cachedTokenCount(ep scheduling.Endpoint) int {
+	tokens, _ := p.cachedTokens(ep)
+	return tokens
 }

pkg/epp/framework/plugins/requestcontrol/dataproducer/p2psource/producer_test.go

@@ -374,22 +374,51 @@ func TestProduce_NearTie_NoHerding(t *testing.T) {
 	assert.Greater(t, busy, 600/6, "busy peer must not be starved: %v", picks)
 }
 
-// A strictly-better-cached peer wins regardless of load: the tie-break
-// applies to equal counts only.
-func TestProduce_HigherCachedWinsOverIdle(t *testing.T) {
+// A peer more than one block ahead wins regardless of load: the queue
+// weighting applies within the one-block band only.
+func TestProduce_BeyondBandWinsOverIdle(t *testing.T) {
 	ctx := utils.NewTestContext(t)
 	p := New("test", Config{MinCachedTokenDelta: 1})
 
 	req := &scheduling.InferenceRequest{RequestID: "req-strict"}
 	eps := []scheduling.Endpoint{
-		endpointWithLoad(p, "pod-a", "10.0.0.1", 4, 20),
+		endpointWithLoad(p, "pod-a", "10.0.0.1", 5, 20),
 		endpointWithLoad(p, "pod-b", "10.0.0.2", 3, 0),
 	}
 	require.NoError(t, p.Produce(ctx, req, eps))
 	best, ok := scheduling.ReadRequestAttribute[*bestMatchPeer](req, p.attrKey())
 	require.True(t, ok)
 	assert.Equal(t, "10.0.0.1:8080", best.hostPort)
-	assert.Equal(t, 4*testBlockSize, best.cachedTokens)
+	assert.Equal(t, 5*testBlockSize, best.cachedTokens)
+}
+
+// A peer one block short of the maximum competes on queue weight: an idle
+// one-block-short peer takes most of the traffic from a deeply-queued
+// maximum, and the stashed count is the chosen peer's own.
+func TestProduce_OneBlockShort_SharesByQueue(t *testing.T) {
+	ctx := utils.NewTestContext(t)
+	p := New("test", Config{MinCachedTokenDelta: 1})
+
+	// Weights: pod-ahead 1/21, pod-short 1.
+	eps := []scheduling.Endpoint{
+		endpointWithLoad(p, "pod-ahead", "10.0.0.1", 4, 20),
+		endpointWithLoad(p, "pod-short", "10.0.0.2", 3, 0),
+	}
+	picks := map[string]int{}
+	for i := 0; i < 400; i++ {
+		req := &scheduling.InferenceRequest{RequestID: fmt.Sprintf("band-%d", i)}
+		require.NoError(t, p.Produce(ctx, req, eps))
+		best, ok := scheduling.ReadRequestAttribute[*bestMatchPeer](req, p.attrKey())
+		require.True(t, ok)
+		picks[best.hostPort]++
+		want := 4 * testBlockSize
+		if best.hostPort == "10.0.0.2:8080" {
+			want = 3 * testBlockSize
+		}
+		assert.Equal(t, want, best.cachedTokens)
+	}
+	assert.Greater(t, picks["10.0.0.2:8080"], picks["10.0.0.1:8080"], "idle one-block-short peer must lead: %v", picks)
+	assert.Greater(t, picks["10.0.0.1:8080"], 0, "max-cached peer must keep some share: %v", picks)
 }
 
 // Peers tied on cached count and queue depth: requests spread across them by
@@ -433,7 +462,7 @@ func TestProduce_NilMetrics_NeutralLoad(t *testing.T) {
 
 	req := &scheduling.InferenceRequest{RequestID: "req-nilmetrics"}
 	eps := []scheduling.Endpoint{
-		endpoint(p, "pod-a", "10.0.0.1", 2),
+		endpoint(p, "pod-a", "10.0.0.1", 1),
 		endpoint(p, "pod-b", "10.0.0.2", 3),
 	}
 	require.NoError(t, p.Produce(ctx, req, eps))