fix(epp): use exact cached blocks for in-flight load#2091
Open
Prasannajaga wants to merge 1 commit into
Open
Conversation
The in-flight load producer incorrectly used the tier-weighted MatchBlocks for literal token load tracking, inflating perceived load when serving from CPU cache. This updates the accounting logic to use the exact CachedBlockCount. Resolves issue llm-d#1292. Signed-off-by: prasanna <prasannajaga9@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The Problem
The router represents prefix-cache matches with
PrefixCacheMatchInfo. Two ofits fields describe different concepts:
MatchBlocksprefix_cache_score.CachedBlockCountThe in-flight load producer currently uses
MatchBlocksas though it were aliteral block count:
That calculation is correct for the approximate prefix producer because its
constructor defaults
CachedBlockCounttoMatchBlocks. It is not correct forthe precise producer when tier weighting makes the two values different.
The narrow correction is:
The existing tail calculation remains unchanged so tokens beyond the indexed
prefix are still counted.
Why We Need This Change (The Impact of the Bug)
Consider a 4,096-token request with 16-token cache blocks. The request contains
256 complete blocks. Assume the full prompt is cached on a pod in CPU cache,
whose default tier weight is
0.8.256 * 0.8 = 204.8MatchBlocks204CachedBlockCount256prefix_cache_score204 / 256 = 0.796875The current in-flight calculation reports:
The literal cached-block calculation reports:
The current result incorrectly treats tier weighting as missing cache data. It
can inflate the current request's
UncachedRequestTokens, inflate the loadcommitted after endpoint selection, and affect later scheduling decisions that
consume accumulated
InFlightLoad.This proposal does not claim that restoring CPU-resident cache has zero latency.
The existing tier-weighted
prefix_cache_scorecontinues to represent relativecache benefit. The change only makes literal token accounting use the literal
cached-block field.
Data Flow
A request contains 4,096 input tokens. Two candidate pods hold the complete
prefix:
MatchBlocks=256andCachedBlockCount=256.MatchBlocks=204andCachedBlockCount=256.The router is configured to use the precise prefix producer for both in-flight
load accounting and predicted-latency cache features.
Step-by-Step Flow
sequenceDiagram autonumber actor Client participant Router participant Token as Token Producer participant Prefix as Precise Prefix Producer participant InFlight as In-Flight Load Producer participant Predict as Predicted Latency Producer participant Score as Scheduling Scorers participant Endpoint as Selected Endpoint Client->>Router: Send 4,096-token request Router->>Token: Produce engine-correlated token IDs Token-->>Router: TokenizedPrompt with 4,096 tokens Router->>Prefix: Look up candidate cache state Prefix-->>Router: Pod A MatchBlocks=256, CachedBlockCount=256 Prefix-->>Router: Pod B MatchBlocks=204, CachedBlockCount=256 Router->>InFlight: Produce current-request load projection Note over InFlight: Old Pod B result: 832 uncached tokens Note over InFlight: Correct Pod B result: 0 uncached tokens InFlight-->>Router: UncachedRequestTokens per endpoint Router->>Predict: Produce latency predictions Note over Predict: prefix_cache_score remains 204/256 for Pod B Note over Predict: Prediction reads accumulated InFlightLoad Predict-->>Router: LatencyPredictionInfo per endpoint Router->>Score: Score candidates and select endpoint Score-->>Router: Selected endpoint Router->>InFlight: PreRequest commits selected request load Router->>Endpoint: Forward request Endpoint-->>Router: Stream response Router->>InFlight: Release committed load Router-->>Client: Return responseProposed Changes
[MODIFY] producer.go
Change the cached-token operand in
uncachedInputTokens:Update the function comment so it describes the literal cached-block contract.
Keep the current fallback, indexed-prefix, tail, and overestimate behavior.
YAML Configuration Comparison
Default (Approximate) Setup
With no explicit producer names, the dependency graph uses the default
approximate prefix-cache producer.
Backward Compatibility: Why This Does Not Affect Existing (Approximate) Deployments
NewPrefixCacheMatchInfoinitializes the literal count from the existing matchcount:
The approximate prefix producer uses this default and does not replace
CachedBlockCount. Therefore, for valid approximate attributes:Switching the token calculation to
CachedBlockCountproduces the same resultfor existing approximate deployments. No configuration field is removed or
renamed, and approximate accounting remains the default.
Opt-in Precise Setup (The Target Fix)
In this setup:
CachedBlockCountfor literal tokenaccounting.
MatchBlocksforprefix_cache_score.InFlightLoadafter arequest is dispatched.
Verification Plan
The focused test matrix should verify:
MatchBlocks != CachedBlockCount.remains unchanged.
Producepublishes literalUncachedRequestTokensfor every candidate.PreRequestcommits the literal uncached input plus any configured outputestimate, and release removes the same stored amount.
The exact regression example is:
Which issue(s) this PR fixes:
partially Fixes #1292
Release note (write
NONEif no user-facing change):