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
| disable=not self.enable_prefix_cache, | ||
| retain_pages=self.retain_pages_for_cache, | ||
| release_pages=self.release_pages_from_cache, | ||
| ) |
There was a problem hiding this comment.
这种page_size 和block_size不对齐的场景什么情况下会出现？
There was a problem hiding this comment.
这里不是运行态允许 page_size 和 block_size 不对齐，而是兼容离线推理的延迟初始化流程。离线 LLMEngine 会先通过 KvCacheManager() 使用默认 block_size=64 创建尚未分配物理 block 的 manager，加载模型后 register_model() 才拿到实际的 runtime.page_size（例如 Qwen 为 128）。因此 _init_blocks() 会在任何物理 block 和有效 Radix 条目创建之前，按最终 page size 重建空的 RadixCache。在线 serving 在构造 manager 时已经传入实际 page size，通常不会进入这个分支。相关初始化时序说明已补充到代码注释中。
| missing_tokens = (num_blocks - self.num_free_blocks) * self.block_size | ||
| self.radix_cache.evict(missing_tokens) | ||
| if self.num_free_blocks < num_blocks: | ||
| return None |
There was a problem hiding this comment.
if self.num_free_blocks < num_blocks: 为什么出现两次。。。
There was a problem hiding this comment.
第一次判断用于发现空间不足并触发 Radix LRU eviction；第二次判断用于验证 eviction 后是否真的释放了足够 block。由于被活动请求锁定/保护的 Radix 路径不能淘汰，evict() 可能只能释放部分空间，因此仍需以 free queue 的实际数量再次确认，不足时返回 None 交给 scheduler 做等待或抢占。为避免看起来像重复判断，已经把复查嵌套到第一次判断中并补充了注释，行为不变。
e63e59c to
ea1d3fe
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: 14 passedpython -m pytest -q tests/test_parallel.py: 9 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 post-rebase workload ran on one Ascend A2/A3 NPU with page size 128. Two trials used opposite backend execution orders to check for order bias:
Deterministic token and step counts below are shown per three-group trial. Timing distributions combine all six groups and 96 requests per backend.
The aggregate worker-step speedup was 1.455x with Hash first and 1.448x with Radix first, so the measured effect was not sensitive to backend execution order.
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. Model loading, compilation, KV allocation, and initialization warmup are excluded. This is an intentionally prefix-heavy workload and should not be interpreted as a universal 1.45x speedup for unrelated prompts.The measurements were collected on 2026-08-03 after rebasing onto
main, using the tree in PR commit5d5a53e. The current Qwen3 accuracy, Qwen3 serving, and Radix CI guards pass on that tree. The benchmark implementation and detailed local report remain outside this PR.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.