feat: add balanceBy tokens option to burst prefix cache producer

Signed-off-by: Ezra Silvera <ezra@il.ibm.com>

Author: ezrasilvera

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

@@ -61,6 +61,7 @@ This producer emits `PrefixCacheMatchInfo`; it does not score. Reuse the
 | `maxPrefixTokensToMatch` | 0 | cap on matched prefix tokens; 0 uses the default block cap. A positive value must be >= `blockSizeTokens`, otherwise it yields zero prefix blocks. |
 | `minColocateBlocks` | 0 | min shared leading blocks for inter-unit co-location and for a single request to gain an affinity; 0 disables both (only identical groups are placed, purely load-balanced) |
 | `maxBatchSize` | 1000 | max requests one window may accumulate; Produce returns an error once reached. -1 = unlimited. |
+| `balanceBy` | `requests` | quantity balanced across replicas within a window: `requests` (by request count) or `tokens` (by prefix-block load, charging each unit's prefix once per replica and discounting the leading blocks already held there). Token balancing spreads long-prompt units so no replica saturates on stacked prefixes while request counts stay even. |
 
 ## Operational notes
 

pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/assign.go

@@ -145,7 +145,7 @@ outer:
 // over the batch (longest-prefix first) so a shared prefix is prefilled once per
 // replica rather than scattered. Identical groups are placed before prefix-sharing
 // singletons so proven same-prompt co-location anchors the layout.
-func assign(entries []*entry, k, minColocateBlocks int) {
+func assign(entries []*entry, k, minColocateBlocks int, byTokens bool) {
 	groups := map[string][]*entry{}
 	order := []string{}
 	var replicas []fwksched.Endpoint
@@ -169,18 +169,16 @@ func assign(entries []*entry, k, minColocateBlocks int) {
 	shared := sharedPrefixKeys(groups, order, minColocateBlocks)
 	var groupUnits, singleUnits []string
 	counts := map[string]int{} // prefix-block count per placed unit, computed once
-	total := 0
 	for _, key := range order {
 		switch {
 		case len(groups[key]) >= 2:
 			groupUnits = append(groupUnits, key)
-			counts[key] = totalBlocks(groups[key][0].hashes)
-			total += len(groups[key])
 		case shared[key]:
 			singleUnits = append(singleUnits, key)
-			counts[key] = totalBlocks(groups[key][0].hashes)
-			total += len(groups[key])
+		default:
+			continue
 		}
+		counts[key] = totalBlocks(groups[key][0].hashes)
 	}
 	if len(groupUnits) == 0 && len(singleUnits) == 0 {
 		return
@@ -191,17 +189,28 @@ func assign(entries []*entry, k, minColocateBlocks int) {
 	sortByPrefixLen(groupUnits, counts)
 	sortByPrefixLen(singleUnits, counts)
 
-	maxShare := (total + len(replicas) - 1) / len(replicas) // ceil: equal samples per replica
+	// total and load carry one unit: request count, or prefix blocks under token
+	// balancing. Token totals sum full per-unit block counts (a loose fair-share
+	// target); the shared-prefix discount is applied to load at placement time.
+	total := 0
+	for key := range counts {
+		if byTokens {
+			total += counts[key]
+		} else {
+			total += len(groups[key])
+		}
+	}
+	maxShare := (total + len(replicas) - 1) / len(replicas) // ceil: fair share per replica
 
 	idx := newBatchIndex()
-	load := map[string]int{} // batch-wide samples assigned per replica
+	load := map[string]int{} // per-replica load in the balancing unit (requests or blocks)
 	// Groups first, then prefix-sharing singletons: same order as before, without
 	// allocating a combined slice.
 	for _, key := range groupUnits {
-		placeGroup(groups[key], replicas, k, minColocateBlocks, maxShare, idx, load)
+		placeGroup(groups[key], replicas, k, minColocateBlocks, maxShare, counts[key], idx, load, byTokens)
 	}
 	for _, key := range singleUnits {
-		placeGroup(groups[key], replicas, k, minColocateBlocks, maxShare, idx, load)
+		placeGroup(groups[key], replicas, k, minColocateBlocks, maxShare, counts[key], idx, load, byTokens)
 	}
 }
 
@@ -217,19 +226,22 @@ func sortByPrefixLen(keys []string, counts map[string]int) {
 // preference bounded by the fair-share cap; remaining samples (when k caps the
 // group) spill to least-loaded replicas. The unit's blocks are then recorded so
 // later units can match against it.
-func placeGroup(members []*entry, replicas []fwksched.Endpoint, k, minColocateBlocks, maxShare int, idx *batchIndex, load map[string]int) {
+func placeGroup(members []*entry, replicas []fwksched.Endpoint, k, minColocateBlocks, maxShare, cost int, idx *batchIndex, load map[string]int, byTokens bool) {
 	if len(replicas) == 0 {
 		return
 	}
 	hashes := members[0].hashes // identical group: any member represents it
 
 	var matches map[string]int
 	preferMatch := minColocateBlocks > 0
-	if preferMatch {
+	// Token balancing also needs the shared-prefix length per replica: a replica
+	// only prefills the blocks it does not already hold, so its load is charged the
+	// remaining cost rather than the full prefix.
+	if preferMatch || byTokens {
 		matches = idx.longestPrefix(hashes)
 	}
 
-	perReplica := map[string]int{} // samples of this group already on a replica
+	perReplica := map[string]int{} // requests of this unit already on a replica
 	i := 0
 	for i < len(members) {
 		target := pickReplica(replicas, perReplica, k, load, matches, minColocateBlocks, maxShare, preferMatch)
@@ -248,8 +260,20 @@ func placeGroup(members []*entry, replicas []fwksched.Endpoint, k, minColocateBl
 		for j := 0; j < run; j++ {
 			members[i+j].assigned = target
 		}
+		firstTouch := perReplica[name] == 0
 		perReplica[name] += run
-		load[name] += run
+		switch {
+		case !byTokens:
+			load[name] += run
+		case firstTouch:
+			// Charge the prefix once per replica this unit lands on, discounting the
+			// leading blocks already held there (prefilled by an earlier request).
+			inc := cost - matches[name]
+			if inc < 0 {
+				inc = 0
+			}
+			load[name] += inc
+		}
 		i += run
 		preferMatch = false // only the primary replica gets the prefix preference
 	}
@@ -322,12 +346,13 @@ func pickLeastLoaded(replicas []fwksched.Endpoint, perReplica map[string]int, k
 	return best
 }
 
-// less reports whether replica a should be preferred over replica b: fewer
-// assigned samples first, then fewer running requests, then lower name. The
-// batch-local assigned-sample count (load) is the primary signal; running
-// requests only break ties within a batch and deliberately differ from the
-// load-aware-scorer's WaitingQueueSize signal, since placement balances work
-// already committed in this window rather than queue depth observed downstream.
+// less reports whether replica a should be preferred over replica b: lower
+// assigned load first, then fewer running requests, then lower name. The
+// batch-local load (assigned requests, or prefix blocks under token balancing) is
+// the primary signal; running requests only break ties within a batch and
+// deliberately differ from the load-aware-scorer's WaitingQueueSize signal, since
+// placement balances work already committed in this window rather than queue depth
+// observed downstream.
 func less(a fwksched.Endpoint, aName string, load map[string]int, b fwksched.Endpoint, bName string) bool {
 	if load[aName] != load[bName] {
 		return load[aName] < load[bName]

pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/assign_test.go

@@ -74,7 +74,7 @@ func TestAssign_UnlimitedColocatesWholeGroup(t *testing.T) {
 	replicas := []fwksched.Endpoint{testEndpoint("pod1"), testEndpoint("pod2"), testEndpoint("pod3"), testEndpoint("pod4")}
 	entries := group(8, []prefixhash.BlockHash{1, 2, 3}, replicas)
 
-	assign(entries, unlimitedPerReplica, 0)
+	assign(entries, unlimitedPerReplica, 0, false)
 
 	for _, e := range entries {
 		assert.Equal(t, "pod1", assignedName(e), "all samples of one group must co-locate when k=-1")
@@ -85,7 +85,7 @@ func TestAssign_CapSpreadsEvenly(t *testing.T) {
 	replicas := []fwksched.Endpoint{testEndpoint("pod1"), testEndpoint("pod2"), testEndpoint("pod3"), testEndpoint("pod4")}
 	entries := group(8, []prefixhash.BlockHash{1, 2, 3}, replicas)
 
-	assign(entries, 2, 0)
+	assign(entries, 2, 0, false)
 
 	assert.Equal(t, map[string]int{"pod1": 2, "pod2": 2, "pod3": 2, "pod4": 2}, counts(entries),
 		"k=2 over 8 samples and 4 replicas must place 2 per replica")
@@ -95,7 +95,7 @@ func TestAssign_CapFillsBeforeSpilling(t *testing.T) {
 	replicas := []fwksched.Endpoint{testEndpoint("pod1"), testEndpoint("pod2"), testEndpoint("pod3"), testEndpoint("pod4")}
 	entries := group(8, []prefixhash.BlockHash{1, 2, 3}, replicas)
 
-	assign(entries, 4, 0)
+	assign(entries, 4, 0, false)
 
 	// k=4 fills one replica to 4 before using the next; only two replicas used.
 	assert.Equal(t, map[string]int{"pod1": 4, "pod2": 4}, counts(entries))
@@ -105,7 +105,7 @@ func TestAssign_SingletonGetsNoAffinity(t *testing.T) {
 	replicas := []fwksched.Endpoint{testEndpoint("pod1"), testEndpoint("pod2")}
 	entries := group(1, []prefixhash.BlockHash{1, 2, 3}, replicas)
 
-	assign(entries, unlimitedPerReplica, 0)
+	assign(entries, unlimitedPerReplica, 0, false)
 
 	assert.Nil(t, entries[0].assigned, "a singleton group has no reuse and must not receive an affinity")
 }
@@ -117,7 +117,7 @@ func TestAssign_EmptyPrefixGetsNoAffinity(t *testing.T) {
 		{hashes: nil, pods: replicas},
 	}
 
-	assign(entries, unlimitedPerReplica, 0)
+	assign(entries, unlimitedPerReplica, 0, false)
 
 	for _, e := range entries {
 		assert.Nil(t, e.assigned, "requests with no prefix must not be grouped")
@@ -130,7 +130,7 @@ func TestAssign_DistinctGroupsSpreadAcrossReplicas(t *testing.T) {
 	groupB := group(8, []prefixhash.BlockHash{9, 9, 9}, replicas)
 	entries := append(append([]*entry{}, groupA...), groupB...)
 
-	assign(entries, unlimitedPerReplica, 0)
+	assign(entries, unlimitedPerReplica, 0, false)
 
 	// Each group co-locates, and the second group lands on a different, less
 	// loaded replica than the first.
@@ -154,7 +154,7 @@ func TestAssign_SharedPrefixFamiliesColocateWithinFairShare(t *testing.T) {
 	y2 := group(2, []prefixhash.BlockHash{7, 8, 6}, replicas)
 	entries := concat(x1, x2, y1, y2)
 
-	assign(entries, unlimitedPerReplica, 2)
+	assign(entries, unlimitedPerReplica, 2, false)
 
 	// Each family co-locates onto one replica, the two families land on
 	// different replicas, and the batch stays balanced (4 samples per replica).
@@ -174,7 +174,7 @@ func TestAssign_FairShareSpreadsPrefixSharingGroups(t *testing.T) {
 	g4 := group(2, []prefixhash.BlockHash{1, 2, 6}, replicas)
 	entries := concat(g1, g2, g3, g4)
 
-	assign(entries, unlimitedPerReplica, 2)
+	assign(entries, unlimitedPerReplica, 2, false)
 
 	assert.Equal(t, map[string]int{"pod1": 2, "pod2": 2, "pod3": 2, "pod4": 2}, counts(entries),
 		"the fair-share cap must spread prefix-sharing groups, not stampede them onto one replica")
@@ -190,7 +190,7 @@ func TestAssign_SingletonAttachesToOverlappingGroup(t *testing.T) {
 	sy := group(1, []prefixhash.BlockHash{7, 8, 9}, replicas)
 	entries := concat(gx, sx, gy, sy)
 
-	assign(entries, unlimitedPerReplica, 2)
+	assign(entries, unlimitedPerReplica, 2, false)
 
 	// The identical group stays whole, and the overlapping singleton attaches to
 	// the replica holding the group it shares a prefix with.
@@ -206,7 +206,7 @@ func TestAssign_LoneSingletonGetsNoAffinity(t *testing.T) {
 	lone := group(1, []prefixhash.BlockHash{5, 5, 5}, replicas) // overlaps nothing
 	entries := concat(g, lone)
 
-	assign(entries, unlimitedPerReplica, 2)
+	assign(entries, unlimitedPerReplica, 2, false)
 
 	assert.Nil(t, lone[0].assigned, "a request overlapping no other must keep no affinity")
 	assert.Equal(t, assignedName(g[0]), assignedName(g[1]), "the identical group still co-locates")
@@ -222,7 +222,7 @@ func TestAssign_OverlappingSingletonsColocateWithoutGroups(t *testing.T) {
 	s4 := group(1, []prefixhash.BlockHash{1, 2, 6}, replicas)
 	entries := concat(s1, s2, s3, s4)
 
-	assign(entries, unlimitedPerReplica, 2)
+	assign(entries, unlimitedPerReplica, 2, false)
 
 	for _, e := range entries {
 		assert.NotNil(t, e.assigned, "overlapping singletons receive an affinity even without identical groups")
@@ -240,7 +240,7 @@ func TestAssign_SharedPrefixBelowThresholdDoesNotColocate(t *testing.T) {
 	y2 := group(2, []prefixhash.BlockHash{7, 8, 6}, replicas)
 	entries := concat(x1, x2, y1, y2)
 
-	assign(entries, unlimitedPerReplica, 3)
+	assign(entries, unlimitedPerReplica, 3, false)
 
 	// The shared prefix falls short of the threshold, so groups are placed purely
 	// by load and a family is not kept together.
@@ -257,21 +257,66 @@ func TestAssign_SingleReplicaPlacesEverything(t *testing.T) {
 	groupB := group(4, []prefixhash.BlockHash{9, 9, 9}, replicas)
 	entries := concat(groupA, groupB)
 
-	assign(entries, unlimitedPerReplica, 2)
+	assign(entries, unlimitedPerReplica, 2, false)
 
 	for _, e := range entries {
 		assert.Equal(t, "pod1", assignedName(e), "with a single replica every request must land on it")
 	}
 }
 
+func TestAssign_BalanceByTokensSpreadsLargePrefix(t *testing.T) {
+	replicas := []fwksched.Endpoint{testEndpoint("pod1"), testEndpoint("pod2")}
+	// One long-prefix unit and two short-prefix units with distinct prefixes.
+	// Balancing by request count puts a short unit on the same replica as the long
+	// one (block-imbalanced); balancing by tokens keeps the long unit alone and
+	// packs the short ones onto the other replica.
+	long := group(2, []prefixhash.BlockHash{1, 2, 3, 4}, replicas)
+	short1 := group(2, []prefixhash.BlockHash{5}, replicas)
+	short2 := group(2, []prefixhash.BlockHash{6}, replicas)
+
+	byReq := concat(long, short1, short2)
+	assign(byReq, unlimitedPerReplica, 0, false)
+	assert.Equal(t, assignedName(long[0]), assignedName(short2[0]),
+		"balancing by requests colocates a short unit with the long-prefix unit")
+
+	long = group(2, []prefixhash.BlockHash{1, 2, 3, 4}, replicas)
+	short1 = group(2, []prefixhash.BlockHash{5}, replicas)
+	short2 = group(2, []prefixhash.BlockHash{6}, replicas)
+	byTok := concat(long, short1, short2)
+	assign(byTok, unlimitedPerReplica, 0, true)
+	assert.NotEqual(t, assignedName(long[0]), assignedName(short1[0]),
+		"balancing by tokens keeps the long-prefix unit off the replica holding the short ones")
+	assert.Equal(t, assignedName(short1[0]), assignedName(short2[0]),
+		"the short units pack onto one replica under token balancing")
+}
+
+func TestAssign_BalanceByTokensDiscountsSharedPrefix(t *testing.T) {
+	replicas := []fwksched.Endpoint{testEndpoint("pod1"), testEndpoint("pod2")}
+	// Three units sharing a 9-block leading prefix. Charging each its full 10-block
+	// prefix would push the third off the shared replica once the fair-share cap is
+	// hit; discounting the already-prefilled 9 blocks keeps all three colocated so
+	// the shared prefix is prefilled once.
+	a1 := group(2, []prefixhash.BlockHash{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, replicas)
+	a2 := group(2, []prefixhash.BlockHash{1, 2, 3, 4, 5, 6, 7, 8, 9, 11}, replicas)
+	a3 := group(2, []prefixhash.BlockHash{1, 2, 3, 4, 5, 6, 7, 8, 9, 12}, replicas)
+	entries := concat(a1, a2, a3)
+
+	assign(entries, unlimitedPerReplica, 9, true)
+
+	assert.Equal(t, assignedName(a1[0]), assignedName(a2[0]),
+		"units sharing a prefix colocate when the shared prefill is not double-charged")
+	assert.Equal(t, assignedName(a2[0]), assignedName(a3[0]),
+		"the shared-prefix discount keeps the third unit on the shared replica")
+}
+
 func TestAssign_OverflowGuardBalancesBeyondCap(t *testing.T) {
 	replicas := []fwksched.Endpoint{testEndpoint("pod1"), testEndpoint("pod2")}
 	// One group of 8 with k=2 over 2 replicas: k*replicas = 4 < 8, so the
 	// per-replica cap cannot hold the whole group. The capLeft overflow guard must
 	// still place every member and keep the batch balanced.
 	entries := group(8, []prefixhash.BlockHash{1, 2, 3}, replicas)
 
-	assign(entries, 2, 0)
+	assign(entries, 2, 0, false)
 
 	for _, e := range entries {
 		assert.NotNil(t, e.assigned, "the overflow guard must still assign every member")

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

@@ -112,6 +112,12 @@ func newDataProducer(_ context.Context, name string, cfg config) (*dataProducer,
 	if cfg.MaxBatchSize != unlimitedBatchSize && cfg.MaxBatchSize < 1 {
 		return nil, fmt.Errorf("invalid configuration: maxBatchSize must be -1 (unlimited) or >= 1 (current value: %d)", cfg.MaxBatchSize)
 	}
+	if cfg.BalanceBy == "" {
+		cfg.BalanceBy = balanceByRequests
+	}
+	if cfg.BalanceBy != balanceByRequests && cfg.BalanceBy != balanceByTokens {
+		return nil, fmt.Errorf("invalid configuration: balanceBy must be %q or %q (current value: %q)", balanceByRequests, balanceByTokens, cfg.BalanceBy)
+	}
 
 	maxBlocks := defaultMaxPrefixBlocks
 	if cfg.MaxPrefixTokensToMatch > 0 {
@@ -192,7 +198,7 @@ func (p *dataProducer) seal(b *batch) {
 	entries := b.entries
 	p.mu.Unlock()
 
-	assign(entries, p.config.MaxPerReplica, p.config.MinColocateBlocks)
+	assign(entries, p.config.MaxPerReplica, p.config.MinColocateBlocks, p.config.BalanceBy == balanceByTokens)
 	close(b.sealed)
 }
 

pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/plugin_test.go

@@ -62,6 +62,9 @@ func TestNew_RejectsInvalidConfig(t *testing.T) {
 
 	_, err = newDataProducer(context.Background(), "burst", config{WindowDurationMs: 100, MaxPerReplica: -1, BlockSizeTokens: 0})
 	assert.Error(t, err, "blockSizeTokens must be > 0")
+
+	_, err = newDataProducer(context.Background(), "burst", config{WindowDurationMs: 100, MaxPerReplica: -1, BlockSizeTokens: 64, BalanceBy: "bogus"})
+	assert.Error(t, err, "balanceBy must be requests or tokens")
 }
 
 func TestProduce_ColocatesIdenticalPromptBurst(t *testing.T) {

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

@@ -51,6 +51,19 @@ const (
 	maxWindowDurationMs = 10000
 )
 
+// balanceMode selects the quantity balanced across replicas within a window.
+type balanceMode string
+
+const (
+	// balanceByRequests balances placement by request count (the default).
+	balanceByRequests balanceMode = "requests"
+	// balanceByTokens balances placement by prefix-block (token) load, charging a
+	// prefix only the blocks a replica must still prefill - blocks already held from
+	// an earlier request sharing that prefix are free. Spreads long-prompt requests
+	// so no replica saturates on stacked prefixes while shorter ones pack together.
+	balanceByTokens balanceMode = "tokens"
+)
+
 // config defines the configuration for the burst prefix cache producer.
 type config struct {
 	// WindowDurationMs is the batch window T in milliseconds. Requests arriving
@@ -78,6 +91,10 @@ type config struct {
 	// is reached, Produce returns an error instead of letting a sustained burst or
 	// a long window grow the batch unbounded. -1 disables the cap.
 	MaxBatchSize int `json:"maxBatchSize"`
+	// BalanceBy selects the quantity balanced across replicas within a window:
+	// "requests" (default) or "tokens". See balanceMode. An empty value defaults
+	// to "requests".
+	BalanceBy balanceMode `json:"balanceBy"`
 }
 
 // defaultConfig provides sensible defaults for the burst prefix cache producer.
@@ -88,6 +105,7 @@ var defaultConfig = config{
 	MaxPrefixTokensToMatch: 0,
 	MinColocateBlocks:      defaultMinColocateBlocks,
 	MaxBatchSize:           defaultMaxBatchSize,
+	BalanceBy:              balanceByRequests,
 }
 
 // entry is one request collected into a batch window.