feat(sessionprefixcache): session-granularity prefix-cache affinity producer (tokenizer-free)#14
Conversation
…producer Add the content-chunking, per-pod chain index, and identity-resolution primitives for a tokenizer-free session-granularity prefix-cache affinity producer, plus table-driven tests covering affinity after seeding, partial and divergent matches, chunk framing, identity precedence, downward usage refinement, and nil guards. The tests fail to build: no producer wires these primitives into the DataProducer/PreRequest/ResponseBody hooks yet, so no PrefixCacheMatchInfo is produced from the content index. Refs llm-d#1980 Signed-off-by: ChethanUK <chethanuk@outlook.com>
Wire the session-granularity prefix-cache producer into the DataProducer, PreRequest, and ResponseBodyProcessor hooks so it publishes PrefixCacheMatchInfo from the content chain index: - Produce scores each candidate pod by its longest cached chain prefix. - PreRequest seeds the served endpoint of every scheduling profile (so P/D-disaggregated prefill nodes gain affinity), deduped, asynchronously. - ResponseBody confirms the prefix the engine actually cached from reported prompt-token usage and trims the over-estimated estimated tail, so the index refines downward rather than only growing. Register the producer with fwkplugin.Register (not as the default owner of PrefixCacheMatchInfoDataKey, which the approximate producer keeps) and document it in the data-producer README. Opt-in and tokenizer-free. Refs llm-d#1980 Signed-off-by: ChethanUK <chethanuk@outlook.com>
|
CodeAnt AI is reviewing your PR. |
Thanks for using CodeAnt! 🎉We're free for open-source projects. if you're enjoying it, help us grow by sharing. Share on X · |
|
Warning Review limit reached
Next review available in: 41 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (9)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
codeant-ai
Bot
added
the
size:XXL
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a new tokenizer-free session-prefix-cache-producer plugin that derives prefix-cache affinity from request content at session granularity. It hashes framed text into chains of fixed-size byte chunks and maintains a per-pod LRU index. Feedback on the implementation highlights a critical compilation error where sync.WaitGroup is incorrectly called with a .Go() method. Additionally, reviewers recommended eliminating a potential memory leak risk by recomputing the chain in PreRequest instead of storing it in p.state, and suggested optimizing the chunking process by using a more efficient UTF-8 boundary scan and reusing the xxhash.Digest instance to reduce heap allocations.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| p.wg.Go(func() { | ||
| for _, srv := range servers { | ||
| p.index.Add(chain, srv, estimated) | ||
| } | ||
| }) |
There was a problem hiding this comment.
The standard library sync.WaitGroup does not have a Go method. This will cause a compilation error. Use p.wg.Add(1) and a standard goroutine with defer p.wg.Done() instead.
p.wg.Add(1)
go func() {
defer p.wg.Done()
for _, srv := range servers {
p.index.Add(chain, srv, estimated)
}
}()| defer p.state.Delete(req.RequestID) | ||
|
|
||
| st, err := fwkplugin.ReadPluginStateKey[*chainState](p.state, req.RequestID, p.stateKey) | ||
| if err != nil || st == nil || len(st.chain) == 0 { | ||
| return | ||
| } |
There was a problem hiding this comment.
Using p.state to store the chunk chain between Produce and PreRequest introduces a memory leak risk. If a request is aborted, timed out, or rejected after Produce but before PreRequest is called, the state entry will never be deleted and will leak in memory. Since buildChain is pure and extremely fast, you can completely eliminate p.state (and its associated complexity/leak risk) by simply recomputing the chain in PreRequest using p.buildChain(req).
chain := p.buildChain(req)
if len(chain) == 0 {
return
}| var chain []uint64 | ||
| var prev uint64 | ||
| for i := 0; i < len(stream) && len(chain) < maxChunks; { | ||
| end := i | ||
| for end < len(stream) && end-i < chunkSize { | ||
| _, size := utf8.DecodeRune(stream[end:]) | ||
| end += size | ||
| } | ||
| if end-i < chunkSize { | ||
| break // trailing partial chunk: dropped | ||
| } | ||
| chunk := stream[i:end] | ||
|
|
||
| var h uint64 | ||
| if len(chain) == 0 { | ||
| h = hashRoot(model, salt, declaredID, chunk) | ||
| } else { | ||
| h = hashNext(prev, chunk) | ||
| } | ||
| chain = append(chain, h) | ||
| prev = h | ||
| i = end | ||
| } |
There was a problem hiding this comment.
This implementation can be significantly optimized for performance and memory efficiency:
- UTF-8 Boundary Scanning: Instead of calling
utf8.DecodeRunein a loop for every single byte of the chunk, you can jump directly toi + chunkSizeand walk forward to skip any UTF-8 continuation bytes (which start with bits10xxxxxx). This reduces the number of operations fromO(chunkSize)toO(1)per chunk. - xxhash Allocation: Instead of allocating a new
xxhash.Digestviaxxhash.New()for every single chunk in the loop, instantiate a single digest outside the loop and calld.Reset()on each iteration. This drastically reduces heap allocations and GC pressure in the hot path.
var chain []uint64
var prev uint64
d := xxhash.New()
for i := 0; i < len(stream) && len(chain) < maxChunks; {
end := i + chunkSize
if end > len(stream) {
break // trailing partial chunk: dropped
}
// Walk forward to avoid splitting a multi-byte UTF-8 rune.
// In UTF-8, continuation bytes start with bits 10xxxxxx (0x80).
for end < len(stream) && (stream[end]&0xC0 == 0x80) {
end++
}
chunk := stream[i:end]
d.Reset()
if len(chain) == 0 {
writeSeeded(d, model)
writeSeeded(d, salt)
writeSeeded(d, declaredID)
_, _ = d.Write(chunk)
} else {
var le [8]byte
binary.LittleEndian.PutUint64(le[:], prev)
_, _ = d.Write(le[:])
_, _ = d.Write(chunk)
}
h := d.Sum64()
chain = append(chain, h)
prev = h
i = end
}
|
CodeAnt AI finished reviewing your PR. |
PreRequest was seeding the per-pod chain index asynchronously via wg.Go. ResponseBody may then confirm a shorter prefix and TrimEstimatedTail. A late async Add(estimated) can re-insert the discarded tail and undo downward refinement (R1-5). Seed synchronously so refine always wins. Signed-off-by: Hermes Agent <hermes@nousresearch.com>
|
🚨 Unsigned commits detected! Please sign your commits. For instructions on how to set up GPG/SSH signing and verify your commits, please see GitHub Documentation. |
2 participants
User description
What type of PR is this?
/kind feature
What this PR does / why we need it:
Adds an opt-in
DataProducer,session-prefix-cache-producer, that maintains prefix-cache affinity at session granularity and feeds the existingprefix-cache-scorerslot viaprefixMatchInfoProducerName.Problem: the two existing producers each pay a cost to find session prefixes —
approx-prefix-cache-producertokenizes at the router and simulates an LRU per engine (bistable on template-sharing traffic);precise-prefix-cache-produceris truthful but consumes per-block KV events and mirrors engine internals. Multi-turn / agentic sessions are byte-prefix extensions turn over turn; deriving identity from that structure gives deterministic affinity with no tokenizer and zero engine changes.Root cause: no
DataProducerpublishedPrefixCacheMatchInfofrom a content-derived index. The scorer already binds any named producer viacfg.PrefixMatchInfoProducerName; only the tokenizing approx producer was registered as the default source.Fix (L1+L2 only):
prompt_cache_key, session headers) seed the chain root only.PreRequestseedsEstimatedon every profile target endpoint;ResponseBodyconfirms from engineusageand trims estimated tail (downward refinement).Producepublishes covered fraction per endpoint as standardPrefixCacheMatchInfo.Out of scope (tracked in umbrella llm-d#1979): L3 truthful extents from engine-labeled events, L4 retention, engine-side steps, full
previous_response_id→root map.Which issue(s) this PR fixes:
Fixes llm-d#1980
How to test:
Local QA evidence (11/11 unit tests PASS with
-race; format green; package lint 0 issues;make build-eppgreen).Release note (write
NONEif no user-facing change):CodeAnt-AI Description
Add a tokenizer-free session prefix-cache producer
What Changed
Impact
✅ Fewer prefix-cache misses across multi-turn sessions✅ Clearer pod choice for repeated prompts✅ No tokenizer required for prefix-cache affinity💡 Usage Guide
Checking Your Pull Request
Every time you make a pull request, our system automatically looks through it. We check for security issues, mistakes in how you're setting up your infrastructure, and common code problems. We do this to make sure your changes are solid and won't cause any trouble later.
Talking to CodeAnt AI
Got a question or need a hand with something in your pull request? You can easily get in touch with CodeAnt AI right here. Just type the following in a comment on your pull request, and replace "Your question here" with whatever you want to ask:
This lets you have a chat with CodeAnt AI about your pull request, making it easier to understand and improve your code.
Example
Preserve Org Learnings with CodeAnt
You can record team preferences so CodeAnt AI applies them in future reviews. Reply directly to the specific CodeAnt AI suggestion (in the same thread) and replace "Your feedback here" with your input:
This helps CodeAnt AI learn and adapt to your team's coding style and standards.
Example
Retrigger review
Ask CodeAnt AI to review the PR again, by typing:
Check Your Repository Health
To analyze the health of your code repository, visit our dashboard at https://app.codeant.ai. This tool helps you identify potential issues and areas for improvement in your codebase, ensuring your repository maintains high standards of code health.