feat: add stage-one radix prefix cache - #99
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughAdds an optional page-aligned Radix prefix-cache backend with CLI and engine wiring, KV ownership tracking, scheduler integration, tests, an NPU benchmark, and documentation. Hash remains the default backend. ChangesRadix Prefix Cache
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant ReplicaEngineCore
participant Scheduler
participant KvCacheManager
participant RadixCache
Client->>ReplicaEngineCore: submit prompt
ReplicaEngineCore->>Scheduler: enqueue Request with model_id
Scheduler->>KvCacheManager: match_radix_prefix(token_ids)
KvCacheManager->>RadixCache: match_prefix(RadixKey)
RadixCache-->>KvCacheManager: matched page slots
Scheduler->>Scheduler: allocate unmatched pages
Scheduler->>KvCacheManager: insert_radix_prefix(completed prefix)
Scheduler-->>Client: generated output
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
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 |
![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 Stage 1 SGLang-style Radix prefix cache backend as an alternative to the default Hash-based prefix cache. It adds the RadixCache implementation (a page-aligned compressed radix tree), a RequestKVPool for logical slot mapping, and updates the Scheduler to support radix matching, incremental caching, and in-batch deduplication (deferring requests to reuse cold shared prefixes). Additionally, CLI options, comprehensive unit tests, and a performance benchmark on the Qwen3-14B NPU path are provided. No review comments were submitted, so there is no feedback to address.
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.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (4)
tests/bench_radix_npu.py (2)
332-335: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd
strict=Trueto the output-equalityzip.If the two backends' group counts ever diverge (e.g. a run terminates early),
zipsilently truncates the comparison instead of surfacing the mismatch, undermining the "Greedy output token IDs identical" correctness check this benchmark relies on.🐛 Proposed fix
outputs_identical = all( left.output_token_ids == right.output_token_ids - for left, right in zip(results[0].groups, results[1].groups) + for left, right in zip(results[0].groups, results[1].groups, strict=True) )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/bench_radix_npu.py` around lines 332 - 335, Update the output-equality comparison that assigns outputs_identical to call zip with strict=True, so differing group counts raise an error instead of silently truncating the correctness check.Source: Linters/SAST tools
134-138: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoffBenchmark reaches into
ReplicaEngineCoreprivate internals.
_pause_background_looptogglescore._running/core._loop_task, and_drive_group(lines 186-187) pushes/pulls directly oncore._input_queue/core._output_queue. These are private attributes of a class defined inasync_engine.py(a prior layer); any refactor there can silently break this benchmark without a public contract signal.♻️ Suggested direction
Expose a small public API on
ReplicaEngineCore(e.g.pause_background_loop()/resume_background_loop()andsubmit_step(command) -> step_output) that the benchmark can call instead of touching_running,_loop_task,_input_queue,_output_queuedirectly.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/bench_radix_npu.py` around lines 134 - 138, Replace the benchmark’s direct access to ReplicaEngineCore private state in _pause_background_loop and _drive_group with a small public API on ReplicaEngineCore, such as pause_background_loop(), resume_background_loop(), and submit_step(command) -> step_output. Update the benchmark to use these methods instead of _running, _loop_task, _input_queue, or _output_queue, while preserving its current pause/resume and step-processing behavior.pypto_serving/serving/sched/scheduler.py (2)
277-319: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDefer path correctness relies on an implicit
num_computed_tokens == 0invariant.The deferral branch calls
_release_radix_request(request), which resets radix state but notnum_computed_tokens. This is only safe because_should_defer_for_radix_producerreturnsFalsewheneverrequest.num_computed_tokens > 0, so the branch is reachable only when_match_radix_requestmatched zero slots. The invariant holds today, but a defensive reset would make it robust against future refactors of the match/defer sequencing.♻️ Optional defensive reset in `_release_radix_request`
def _release_radix_request(self, request: Request) -> None: pages = self.request_kv_pool.free(request.request_id) if pages: self.kv_cache_manager.release_pages_from_request(pages) if request.last_node is not None: self.kv_cache_manager.radix_cache.dec_lock_ref(request.last_node) request.cached_block_ids = [] request.allocated_block_ids = [] request.prefix_indices = [] request.last_node = None request.cache_protected_len = 0 request.kv_committed_len = 0 request.num_blocks_cached = 0 + request.num_computed_tokens = 0🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pypto_serving/serving/sched/scheduler.py` around lines 277 - 319, Ensure the radix-cache deferral path preserves a zero computed-token state by resetting request.num_computed_tokens when releasing a deferred request. Update _release_radix_request, or the defer branch around _should_defer_for_radix_producer, so future changes to matching order cannot retain stale computed-token progress.
606-623: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueIn-batch dedup scans all running producers per waiting request.
_should_defer_for_radix_produceris O(running × page_size) per waiting request due to re-tuplingproducer.all_token_ids[:page_size]for each candidate on every schedule pass. It's bounded bymax_num_running_reqsand page size, so it's fine at current scales, but if page_size or batch width grows this becomes a hot-path cost. Consider precomputing/caching each running request's first-page tuple if profiling later shows scheduler CPU pressure.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pypto_serving/serving/sched/scheduler.py` around lines 606 - 623, The current implementation repeatedly constructs each producer’s first-page tuple inside _should_defer_for_radix_producer; if profiling identifies scheduler CPU pressure, cache or precompute the first-page identity for running requests and reuse it during comparisons, while preserving the existing namespace, request, token-count, and page-size checks.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@pypto_serving/serving/sched/scheduler.py`:
- Around line 277-319: Ensure the radix-cache deferral path preserves a zero
computed-token state by resetting request.num_computed_tokens when releasing a
deferred request. Update _release_radix_request, or the defer branch around
_should_defer_for_radix_producer, so future changes to matching order cannot
retain stale computed-token progress.
- Around line 606-623: The current implementation repeatedly constructs each
producer’s first-page tuple inside _should_defer_for_radix_producer; if
profiling identifies scheduler CPU pressure, cache or precompute the first-page
identity for running requests and reuse it during comparisons, while preserving
the existing namespace, request, token-count, and page-size checks.
In `@tests/bench_radix_npu.py`:
- Around line 332-335: Update the output-equality comparison that assigns
outputs_identical to call zip with strict=True, so differing group counts raise
an error instead of silently truncating the correctness check.
- Around line 134-138: Replace the benchmark’s direct access to
ReplicaEngineCore private state in _pause_background_loop and _drive_group with
a small public API on ReplicaEngineCore, such as pause_background_loop(),
resume_background_loop(), and submit_step(command) -> step_output. Update the
benchmark to use these methods instead of _running, _loop_task, _input_queue, or
_output_queue, while preserving its current pause/resume and step-processing
behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 7f1a4c6a-9dda-4139-b99a-f5f6e35cdf23
📒 Files selected for processing (13)
README.mddocs/dev/radix_npu_benchmark.mdpypto_serving/cli/main.pypypto_serving/serving/engine/async_engine.pypypto_serving/serving/memory/kv_cache.pypypto_serving/serving/memory/prefix_cache/__init__.pypypto_serving/serving/memory/prefix_cache/base.pypypto_serving/serving/memory/prefix_cache/radix.pypypto_serving/serving/memory/request_kv_pool.pypypto_serving/serving/sched/scheduler.pytests/bench_radix_npu.pytests/test_qwen3_serving.pytests/test_radix_cache.py
ef5aef8 to
df87896
Compare
|
把你的 radix tree加到ci里看护一下吧 |
df87896 to
80e668d
Compare
已添加 CI 看护。现在 CI 会执行 tests/test_radix_cache.py，覆盖前缀匹配、插入与拆分、引用计数、LRU 淘汰、重复页面收敛及请求清理等 Radix Tree 核心行为，共 11 个测试用例。相关修改已更新到当前 PR。 |
80e668d to
fe25f87
Compare
2 participants
Summary
hashas the defaultunit-testsjobCloses #38.
Behavior
The existing behavior remains unchanged unless the new backend is selected:
hashremains the CLI and engine default. The radix backend operates entirely in the serving scheduler and KV metadata layer; it continues to use the existing block table, slot mapping, paged KV tensors, and Qwen attention kernels.Radix entries are namespaced by model ID and contain only complete KV pages. A request retains matched pages and locks the matched path while active. Completed page-aligned prompt and generated-token prefixes are published incrementally. When duplicate prefixes were computed concurrently, requests converge on the canonical radix pages and release redundant physical pages.
For a cold group of requests sharing the first full page, one request is scheduled as the prefix producer. Other consumers are deferred until the producer commits the shared prefix, then they reuse its canonical KV pages and prefill only their unmatched suffixes.
Validation
python -m pytest -q tests/test_radix_cache.py: 11 passedpython -m pytest -q tests/test_batching.py tests/test_parallel.py: 30 passedgit diff --check: passedQwen3-14B NPU Hash versus Radix comparison
A separate local benchmark was used to quantify the behavior without adding benchmark-only scripts or generated reports to this PR. The workload ran on one Ascend A2/A3 NPU with page size 128:
The Hash backend sees an empty cache at the beginning of each cold group and schedules all 16 prompts, computing
16 * 392 = 6,272prefill tokens. The Radix backend schedules one prefix producer, then the remaining 15 requests reuse its three committed pages and compute only their unique suffixes:392 + 15 * 8 = 512prefill tokens per group.Worker step wall timestarts before the main process submits a step and ends when it receives the worker result. It includes IPC, worker-side input preparation, NPU dispatch and execution, sampling, and result IPC; it is not a pure device-kernel measurement. This is an intentionally prefix-heavy workload and should not be interpreted as a universal 5x speedup for unrelated prompts.The measurements were collected on 2026-07-14 before rebasing onto the latest
main. A post-rebase local NPU rerun is currently blocked during Qwen kernel compilation by a local PyPTO parser/version incompatibility (hidden_statestype reassignment), before either prefix-cache backend executes. The default Hash control fails at the same compile point; CPU scheduler/cache regression coverage passes on the rebased branch.Scope
This Stage 1 change establishes the core radix tree, page ownership, scheduler integration, lifecycle cleanup, cold in-batch reuse, and tests. The standalone NPU benchmark implementation and detailed report remain in a separate local commit and are intentionally not included in this PR because they are validation-only artifacts. More advanced SGLang scheduling policies, richer cache namespaces, session semantics, and optional partial-page behavior remain follow-up work.