Skip to content

fix(epp): use exact cached blocks for in-flight load#2091

Open
Prasannajaga wants to merge 1 commit into
llm-d:mainfrom
Prasannajaga:fix-issue-1292
Open

fix(epp): use exact cached blocks for in-flight load#2091
Prasannajaga wants to merge 1 commit into
llm-d:mainfrom
Prasannajaga:fix-issue-1292

Conversation

@Prasannajaga

Copy link
Copy Markdown

The Problem

The router represents prefix-cache matches with PrefixCacheMatchInfo. Two of
its fields describe different concepts:

Field Meaning Intended usefix(epp): use exact cached blocks for in-flight load
MatchBlocks Tier-weighted cache-match score. A CPU-resident block may contribute less than one block. Cache ranking and the existing predicted-latency prefix_cache_score.
CachedBlockCount Literal number of contiguous cached blocks, independent of storage tier. Converting cached blocks into cached or uncached token counts.

The in-flight load producer currently uses MatchBlocks as though it were a
literal block count:

uncached indexed tokens =
    (TotalBlocks - MatchBlocks) * BlockSizeTokens

That calculation is correct for the approximate prefix producer because its
constructor defaults CachedBlockCount to MatchBlocks. It is not correct for
the precise producer when tier weighting makes the two values different.

The narrow correction is:

uncached indexed tokens =
    (TotalBlocks - CachedBlockCount) * BlockSizeTokens

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.

Value Result
Raw weighted score 256 * 0.8 = 204.8
Published MatchBlocks 204
Published CachedBlockCount 256
Existing prefix_cache_score 204 / 256 = 0.796875

The current in-flight calculation reports:

(256 - 204) * 16 = 832 uncached input tokens

The literal cached-block calculation reports:

(256 - 256) * 16 = 0 uncached input tokens

The current result incorrectly treats tier weighting as missing cache data. It
can inflate the current request's UncachedRequestTokens, inflate the load
committed 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_score continues to represent relative
cache 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:

  • Pod A holds 256 blocks in GPU cache.
  • Pod B holds 256 blocks in CPU cache.
  • Pod A publishes MatchBlocks=256 and CachedBlockCount=256.
  • Pod B publishes MatchBlocks=204 and CachedBlockCount=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 response
Loading

Proposed Changes

[MODIFY] producer.go

Change the cached-token operand in uncachedInputTokens:

- matched := int64(info.MatchBlocks()) * blockSize
+ cached := int64(info.CachedBlockCount()) * blockSize
  indexed := int64(info.TotalBlocks()) * blockSize

- uncachedIndexed := indexed - matched
+ uncachedIndexed := indexed - cached

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

plugins:
  - type: token-producer

  - type: inflight-load-producer

  - type: predicted-latency-producer

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

NewPrefixCacheMatchInfo initializes the literal count from the existing match
count:

func NewPrefixCacheMatchInfo(matchBlocks int, totalBlocks int, blockSizeTokens int) *PrefixCacheMatchInfo {
 return &PrefixCacheMatchInfo{
  matchBlocks:      matchBlocks,
  totalBlocks:      totalBlocks,
  blockSizeTokens:  blockSizeTokens,
  cachedBlockCount: matchBlocks,
 }
}

The approximate prefix producer uses this default and does not replace
CachedBlockCount. Therefore, for valid approximate attributes:

CachedBlockCount == MatchBlocks

Switching the token calculation to CachedBlockCount produces the same result
for existing approximate deployments. No configuration field is removed or
renamed, and approximate accounting remains the default.

Opt-in Precise Setup (The Target Fix)

plugins:
  - type: token-producer
    parameters:
      modelName: ${MODEL_NAME}
      vllm:
        url: http://localhost:8000

  - name: precise-prefix
    type: precise-prefix-cache-producer
    parameters:
      # Existing KV index and event-source settings belong here.

  - name: precise-inflight
    type: inflight-load-producer
    parameters:
      prefixMatchInfoProducerName: precise-prefix

  - name: predicted-latency
    type: predicted-latency-producer
    parameters:
      prefixMatchInfoProducerName: precise-prefix
      inFlightLoadProducerName: precise-inflight

In this setup:

  • The in-flight producer uses precise CachedBlockCount for literal token
    accounting.
  • The predicted-latency producer uses precise, tier-weighted MatchBlocks for
    prefix_cache_score.
  • Later predictions observe the corrected accumulated InFlightLoad after a
    request is dispatched.

Verification Plan

The focused test matrix should verify:

  1. A nil endpoint returns the non-negative full input length.
  2. Missing or incorrectly typed prefix information falls back to the full input.
  3. A fully cached precise prefix returns zero uncached input when
    MatchBlocks != CachedBlockCount.
  4. A partially cached precise prefix counts only the literal uncached blocks.
  5. A prompt longer than the indexed range includes its unindexed tail.
  6. A partial final block is included in the tail.
  7. A prompt shorter than one cache block is treated as an unindexed tail.
  8. Approximate-style attributes preserve their existing result.
  9. An invalid block size falls back to the full input.
  10. A cached count above the indexed total cannot produce a negative result.
  11. Existing precedence between indexed block data and a shorter token estimate
    remains unchanged.
  12. Produce publishes literal UncachedRequestTokens for every candidate.
  13. PreRequest commits the literal uncached input plus any configured output
    estimate, and release removes the same stored amount.

The exact regression example is:

inputTokens = 4096
MatchBlocks = 204
CachedBlockCount = 256
TotalBlocks = 256
BlockSizeTokens = 16
expected uncached input = 0

Which issue(s) this PR fixes:

partially Fixes #1292

Release note (write NONE if no user-facing change):

NONE

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>
@Prasannajaga
Prasannajaga requested a review from a team as a code owner July 20, 2026 19:55
@Prasannajaga
Prasannajaga requested review from liu-cong and vMaroon July 20, 2026 19:55
@github-actions github-actions Bot added the size/L Denotes a PR that changes 100-499 lines, ignoring generated files. label Jul 20, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

size/L Denotes a PR that changes 100-499 lines, ignoring generated files.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Evaluate predicted latency with precise tokens accounting

1 participant