perf: speed up KV cache block tracking - #1428
zhangjianning-zjn wants to merge 1 commit into
Conversation
| "//rtp_llm/models_py/bindings/cuda/ops:cuda_impl", | ||
| "//rtp_llm/models_py/bindings/core:exec_ops_test_lib", | ||
| "@local_config_cuda//cuda:cuda_headers", | ||
| "@local_config_cuda//cuda:cudart", |
There was a problem hiding this comment.
[P2] block_cache_test_deps 移除 cuda13_torch_link_deps 可能造成 cuda13_x86 链接失败
- 行原文：block_cache_test_deps 由 ["//rtp_llm/cpp/testing:device_test_utils", "//rtp_llm/models_py/bindings/cuda/ops:cuda_impl", "//rtp_llm/models_py/bindings/core:exec_ops_test_lib", "//rtp_llm/cpp/config:config_modules", "//rtp_llm/cpp/config:model_config", "//rtp_llm/cpp/cache", ...] 缩减为仅 ["//rtp_llm/cpp/cache:block_pool", gtest, gtest_main]。本次移除范围远大于已记录的 cuda13_torch_link_deps：device_test_utils、config_modules、model_config、cuda_impl、exec_ops_test_lib、cuda_headers、cudart、torch_deps 全部被删。若 block_cache_test.cc 仍 include BlockCache.h（不在 block_pool 目标可见 srcs 中）或使用 device_test_utils/config 模块，则编译或链接失败。
建议：确认 block_cache_test / shared_block_cache_test 在 using_cuda13_x86 平台能否成功链接；若需要，在 block_cache_test_deps 中保留 + cuda13_torch_link_deps（或至少 torch_deps()），避免 cuda13 平台链接缺 nvshmem 符号。
评审版本：6290bd132888
| std::scoped_lock lock(ref_mu_, free_mu_); | ||
| request_ref_counter_.decrementRefCounter(block_ids); | ||
| req_con_ref_counter_.decrementRefCounter(block_ids); | ||
| auto free_indices = req_con_ref_counter_.decrementRefCounterWithFreeInfo(block_ids); |
There was a problem hiding this comment.
[P2] 释放条件从五个计数器全零缩减为仅检查 req_con 与 block_cache，依赖未文档化的计数器不变量，存在提前释放/use-after-free 风险
旧 tryFreeBlocks 在每次 free 调用后直接检查 req_con==0 && block_cache==0，与计数器如何归零无关。新逻辑仅将 decrementRefCounterWithFreeInfo 返回的 free_indices（本次调用中恰好归零的块）送入 freeBlocks。若 req_con_ref_counter_/block_cache_ref_counter_ 在其他路径（BlockRefCounter 现作为公开库 block_ref_counter 暴露，decrementRefCounter 普通版本仍公开可用）经普通 decrementRefCounter 归零，且另一计数器当时已为 0，则不会产生 free_indices，块永久泄漏。补丁内三个 free 函数均使用 WithFreeInfo，故本分片内无触发，风险在于补丁外调用方或后续改动。
建议：在 BlockRefCounter/BlockPool 处补充不变量注释或断言（assert(req_con_ref_counter_==request+connector)），或在 freeBlocks 前仍校验 request_ref_counter_/connector_ref_counter_ 为 0；至少补充覆盖 requestReference 与 connectorReference 混合引用路径的测试。
评审版本：6290bd132888
PR #1428 评审：未发现阻塞项，等待人工审核
无阻塞项非阻塞发现
已确认修复
自动代码评审 · head |
rtp-llm-review-bot
left a comment
rtp-llm-review-bot
left a comment
There was a problem hiding this comment.
lgtm ready to ci
评审版本：a046e0a01aebca90d9afeffa057fb34b6b12fffe
a046e0a to
6290bd1
Compare
| freeBlocks(free_indices); | ||
| } | ||
|
|
||
| // Must be called with ref_mu_ and free_mu_ held. |
There was a problem hiding this comment.
[P2] freeBlocks 移除内部双计数器校验后，释放安全完全依赖调用方预先过滤，契约未在注释中记录
补丁 - 行显示 tryFreeBlocks 原实现为 for 循环内检查 req_con_ref_counter_==0 && block_cache_ref_counter_==0 才 insert；+ 行改为 free_block_ids_.insert(block_ids.begin(), block_ids.end()) 无条件插入。校验被拆散到三个调用方的 vectorRemoveIf 过滤（requestFree/connectorFree 过滤 block_cache，blockCacheFree 过滤 req_con）。方法注释仅写 "Must be called with ref_mu_ and free_mu_ held"，未记录"调用方必须预先过滤"这一契约。当前三个调用方过滤均正确，但未来新增调用方或改动过滤条件时易引入提前释放/use-after-free。
建议：在 freeBlocks 注释中明确调用方必须先过滤仍被引用的 block；或在 freeBlocks 内保留双计数器校验作为防御。
评审版本：6290bd132888
|
|
||
| std::vector<int> decrementRefCounterWithFreeInfo(const std::vector<int>& block_indices) { | ||
| auto free_block = decrementRefCounterImpl<true>(block_indices); | ||
| return free_block; |
There was a problem hiding this comment.
[P3] decrementRefCounterWithFreeInfo 冗余中间变量且命名误导（free_block 实为 vector）
补丁 + 行：std::vector decrementRefCounterWithFreeInfo(...) { auto free_block = decrementRefCounterImpl(block_indices); return free_block; }。free_block 是单数命名却持有多个 block；可直接 return decrementRefCounterImpl(block_indices)。此外 decrementRefCounterImpl 仍构造并返回空 vector 后被丢弃，有微小开销。
建议：改为直接 return decrementRefCounterImpl(block_indices)，变量名改为 free_blocks； 分支可考虑返回 void 特化。
评审版本：6290bd132888
| "RTP_LLM_PIN_HOST_BLOCK_POOL": "0", | ||
| }, | ||
| # Host allocation still links CUDA driver libraries through block_pool. | ||
| exec_properties = {'gpu':'A10'}, |
There was a problem hiding this comment.
[P3] block_pool_ref_test 是纯 HOST 分配测试却硬编码 exec_properties={'gpu':'A10'} 强制要求 GPU
新增的 block_pool_ref_test（第 181 行起）其 exec_properties = {'gpu':'A10'}（第 193 行）硬编码 CUDA GPU，未使用 select。同文件其他跨平台测试均按平台选择：block_pool_device_malloc_test（第 273-275 行）与 hybrid_pool_config_creator_test（第 75-76 行）都使用 select({"@//:using_rocm": {"gpu":"MI308X-ROCM7"}, ...})。而该测试依赖的 //rtp_llm/cpp/cache:block_pool 库在 @//:using_rocm 分支下链接 hip（见 cache/BUILD 中 block_pool 的 select），因此 ROCm 构建下该测试二进制链接的是 ROCm 运行时，却请求 A10（CUDA GPU），会导致 ROCm CI 无法调度/运行该测试。
建议：该测试仅验证 BlockPool 的引用计数释放语义，可改为直接依赖新拆出的 block_ref_counter 目标（不链接 CUDA），或将 BlockPoolRefTest 的引用计数断言下沉到 utils/test 的 block_ref_counter_test（已存在且无需 GPU），从而去除对 GPU 的硬依赖；若确因 block_pool 链接 CUDA 无法避免，建议至少将 exec_properties 的 GPU 型号改为与仓库其他 host 测试一致的宽松配置而非硬编码 A10。
评审版本：12ee4c81caba
rtp-llm-review-bot
left a comment
rtp-llm-review-bot
left a comment
There was a problem hiding this comment.
lgtm ready to ci
评审版本：6290bd132888f9f512be296e783b6c1311345f72
|
internal source has been updated, please review the changes! |
6290bd1 to
63f8468
Compare
| "@local_config_cuda//cuda:cudart", | ||
| ] + torch_deps() + cuda13_torch_link_deps | ||
| ] | ||
|
|
There was a problem hiding this comment.
[P2] block_cache_test_deps 移除 cuda13_torch_link_deps 可能造成 cuda13_x86 链接失败
补丁 - 行显示 block_cache_test_deps 移除了 "//rtp_llm/cpp/testing:device_test_utils"、"//rtp_llm/models_py/bindings/cuda/ops:cuda_impl"、"//rtp_llm/models_py/bindings/core:exec_ops_test_lib"、"//rtp_llm/cpp/config:config_modules"、"//rtp_llm/cpp/config:model_config"、"@local_config_cuda//cuda:cuda_headers"、"@local_config_cuda//cuda:cudart" 及 torch_deps()，仅保留 //rtp_llm/cpp/cache:block_pool + gtest。这与台账仅针对 cuda13_torch_link_deps 的发现不同，是更广的依赖削减。若 BlockCacheTest.cc 仍引用 device_test_utils/torch/config/CUDA 符号，将导致编译或链接失败，且不限于 cuda13_x86。
建议：确认 block_cache_test / shared_block_cache_test 在 using_cuda13_x86 平台能否成功链接；若需要，在 block_cache_test_deps 中保留 + cuda13_torch_link_deps（或至少 torch_deps()），避免 cuda13 平台链接缺 nvshmem 符号。
评审版本：12ee4c81caba
| } | ||
|
|
||
| private: | ||
| template<bool with_free_info> |
There was a problem hiding this comment.
[P3] decrementRefCounterWithFreeInfo 冗余中间变量且命名误导（free_block 实为 vector）
decrementRefCounter 现调用 decrementRefCounterImpl，该模板体无条件声明 std::vector<int> free_blocks; 并 return free_blocks;。对 实例化，free_blocks 恒为空且返回值被上层丢弃，但仍每次构造一个 vector 对象（无堆分配，开销微小）。改动前 decrementRefCounter 直接内联循环、不构造任何 vector。
建议：改为直接 return decrementRefCounterImpl(block_indices)，变量名改为 free_blocks； 分支可考虑返回 void 特化。
评审版本：12ee4c81caba
| copts = test_copts, | ||
| deps = google_test_deps + [ | ||
| "//rtp_llm/cpp/cache:block_ref_counter", | ||
| "//rtp_llm/cpp/utils:core_utils", |
There was a problem hiding this comment.
[P3] block_ref_counter 库定义在 cache 层，其单测却放在 utils/test，形成 utils/test → cache 的反向分层依赖
本轮新增 cc_test block_ref_counter_test 位于 rtp_llm/cpp/utils/test/BUILD，deps 含 "//rtp_llm/cpp/cache:block_ref_counter",；而该库目标 block_ref_counter 定义在 rtp_llm/cpp/cache/BUILD。utils 层测试反向依赖 cache 层，破坏了自下而上的依赖方向（cache 依赖 utils，而非 utils 测试依赖 cache）。
建议：将 block_ref_counter_test 移至 rtp_llm/cpp/cache/test 下，或将 block_ref_counter 库目标迁入 rtp_llm/cpp/utils，使测试与其被测目标处于同一层。
评审版本：63f8468cedc8
rtp-llm-review-bot
left a comment
rtp-llm-review-bot
left a comment
There was a problem hiding this comment.
lgtm ready to ci
评审版本：63f8468cedc8a9b96b7139c1247601b0eea91266
63f8468 to
12ee4c8
Compare
| block_pool_->requestFree(request_owned_blocks); | ||
| EXPECT_EQ(block_pool_->freeBlocksNum(), total_blocks_); | ||
|
|
||
| // A block is reusable only after its final connector reference is released. |
There was a problem hiding this comment.
[P2] requestFree 的 block_cache 过滤分支缺少直接测试覆盖
新增测试覆盖了 ConnectorFree 在 request 持有时不释放、BlockCacheFree 在 request 持有时不释放，以及 connector/cache 双重引用后才释放；但 requestFree 中新增的 vectorRemoveIf 对 block_cache_ref_counter_ 的过滤（BlockPool.cc:527-529，即 request 释放时 block_cache 仍持有则不移入 free_block_ids_）没有对称测试：未验证先 blockCacheReference 再 requestFree 时块不被释放、随后 blockCacheFree 才释放。
建议：补充一个测试：malloc 后 blockCacheReference，requestFree 断言 freeBlocksNum 不变且 blockCacheRefBlocksNum 不变，再 blockCacheFree 断言恢复 total。
评审版本：12ee4c81caba
rtp-llm-review-bot
left a comment
rtp-llm-review-bot
left a comment
There was a problem hiding this comment.
lgtm ready to ci
评审版本：12ee4c81caba601b3ce572808ed4dbed42503891
|
internal source has been updated, please review the changes! |
4 similar comments
|
internal source has been updated, please review the changes! |
|
internal source has been updated, please review the changes! |
|
internal source has been updated, please review the changes! |
|
internal source has been updated, please review the changes! |
AI Code Review - PR #1428Status: NEEDS REBASE 该 PR 当前无法 rebase（与目标分支存在冲突 / dirty），已跳过本次自动代码审查。请 rebase 到最新目标分支并解决冲突后重新推送，审查会在新 head 上自动重新触发。 |
|
closed in favor of #1447 |
4 participants
Summary
Validation
All related tests pass under local ROCm (config=rocm):
Total: 66/66 tests passed.