Prhub

#32710 [Radix Cache] Add Rust TreeCore backend with shared parity tests

原始 PR 作者 Jialin 合并时间 2026-09-01 00:26 文件变更 72 提交数 51 评论 11 代码增减 +39974 / -397

执行摘要

新增 Rust TreeCore 缓存后端与共享一致性测试套件

PR body 明确动机:'Add an in-tree Rust implementation of UnifiedTreeCoreInterface for Unified Radix Cache while retaining Python cache orchestration, pools, and action application. The same cache-level unit suite now exercises both Python and Rust TreeCore implementations, exposing parity gaps through shared assertions rather than backend-specific test copies.' 即引入 Rust 后端以获得无 Python 节点对象遍历的高性能路径,同时用共享测试套件保证与 Python 实现的语义一致性。此外,issue 讨论中 ishandhanani 与 dreamtalen 提出希望该 Rust 实现可作为独立 crate 用于 DynoSim 纯 Rust 仿真器,降低 SGLang 与其仿真之间的语义漂移。

值得精读的架构级 PR。重点学习:1)FFI 边界设计——owned snapshot、tagged tuple 编解码、禁止迭代器跨 Rust 边界、inspection binding 与生产 binding 严格隔离(wheel 内校验无 inspect_* 方法);2)双后端共享测试套件的一致性验证策略,比后端各自复制测试更能暴露语义差异;3)fail-fast 特性管理——不支持的组合在构造期显式拒绝并附 TODO(Jialin),避免不完整语义静默运行;4)TreeComponent trait 抽象让 Python 语义可以逐方法移植,每个方法都带 Python 参考实现注释,是跨语言移植的良好范例。

讨论亮点

review 中的核心讨论集中在四个方面:

  1. 独立 crate 用于仿真:ishandhanani 提出 'being able to use this rust radix tree in simulation studies without needing to pull in libtorch or other deps';dreamtalen 补充 DynoSim 是纯 Rust 离散事件仿真器,'If the Rust TreeCore in this PR could be made available as a standalone crate, we could use it as DynoSim's radix tree. The main benefit would be reducing semantic drift between SGLang and its simulation.' 该需求未在本 PR 落地,作为后续协作方向。

  2. 双后端等价性测试:ispobock 建议 'drive the same op sequence through both and assert equal device_indices, LRU eviction order, and CacheActions',最终以共享 UnifiedRadixCache 测试套件 + inspector 快照方式实现,而非后端各自复制测试。

  3. 重复 backup 缺陷:alphabetc1 指出 'A double-backup issue' 并附参考修复,已通过独立 commit 'fix: prevent duplicate full kv backups' 修复。

  4. CI 构建竞态:ispobock 定位到 ci_install_dependency.sh 共享 CARGO_TARGET_DIR,CUDA job 的 rm -rf 会打断同机 Rust 编译,最终采用 'building in a per run target dir under RUNNER_TEMP' 解决。

实现拆解

实现按 5 个步骤推进:

  1. 新增 Rust 核心 crate 与 PyO3 绑定:新增 rust/mem-cache crate,含 node.rs(Node/NodeArena/KeyNamespace/ValueState 槽位)、unified_lru_list.rs(哨兵单元双向链表)、unified_tree_core.rs(match/insert/evict 主流程);通过 PyO3 暴露 RustUnifiedTreeCoreBinding 等绑定类,Python 侧 extension.py 负责按环境变量选择加载。

  2. Python adapter 与注册集成:新增 adapter.py 的 RustUnifiedTreeCore 实现 UnifiedTreeCoreInterface,核心工作是 tagged tuple 编解码(cache_action_from_tagged/_kv_event_from_tagged/_transfer*_binding 等),Rust 只回传 (tag, ...) 元组和 owned 数组,不跨 FFI 传迭代器;构造器对 session radix cache、C128、自定义组件等不支持配置直接抛 ValueError。

  3. 组件驱动抽象与语义移植:components/mod.rs 定义 TreeComponent trait(match validator、LRU refresh、节点分裂重分配、逐出、锁路径、insert 提交),full.rs/swa.rs/mamba.rs 各自实现组件语义;随后按 27 个独立 commit 分批移植 #33091(分配感知逐出)、#36317(Full-only load-back 所有权)、#31479(KV 事件合并)、#31902(拒绝 unsafe prefetch-refill)、#33639(增量 Mamba backup)、#34808(DCP 感知 checkpoint grid)、#32228(部分 lock replay)、#30827(cache-salt 命名空间)、#37091(adopted insert ranges)等历史行为。

  4. 边界与失败语义:typed boundary errors、stale-handle 操作抛 KeyError 且不污染核心、poisoned-core fail-closed;生产 binding 不暴露 inspect_* 方法,测试 inspector 独立成模块且不进 wheel。

  5. 测试/构建/CI 配套:共享 UnifiedRadixCache 套件 2436 例双后端通过,Rust 原生测试 831 例、adapter 集成测试 100 例;setup.py 接入 cargo workspace 元数据与扩展构建;prepare_sglang_wheel.py 新增 auditwheel repair 后的 smoke test(校验加载路径、类名、无 inspect 方法);CI 新增 native mem-cache 与 release-wheel 覆盖,并用 per-run target dir 解决 CARGO_TARGET_DIR 并发竞态。

文件 模块 状态 重要度
python/sglang/srt/mem_cache/rust_tree_core/adapter.py 适配层 added 9.17
rust/mem-cache/src/components/mod.rs 组件框架 added 9.08
rust/mem-cache/src/components/mamba.rs Mamba 组件 added 8.89
rust/mem-cache/src/components/full.rs FULL 组件 added 8.89
rust/mem-cache/src/components/swa.rs SWA 组件 added 8.89
rust/mem-cache/src/node.rs 节点结构 added 8.89
python/sglang/srt/mem_cache/buffer_mode/pipeline.py 备份管线 modified 8.69
scripts/release/prepare_sglang_wheel.py 打包脚本 added 8.69
test/registered/unit/mem_cache/test_rust_tree_core_integration.py 集成测试 added 7.97

关键符号

RustUnifiedTreeCore.__init__ _cache_action_from_tagged _kv_event_from_tagged _insert_step_from_binding _match_result_from_binding TreeComponent::create_match_validator TreeComponent::refresh_lru TreeComponent::commit_insert_component_data TreeComponent::evict_component TreeComponent::redistribute_on_node_split FullComponent::finalize_match_result_in_tree_core FullComponent::evict_device_next_node MambaComponent::commit_insert_component_data MambaComponent::evict_excess_path_states SwaComponent::commit_prefetch_ SwaComponent::maybe_split_leaf_for_swa_lock_ KeyNamespaceRef::new UnifiedLRUList::reset_node_and_window_ancestors_mru snapshot_buffer_backup validate_buffer_backup dfs_weight_order

分析完成后,这里会展示 LLM 生成的相对完整源码片段和详细注释。

评论区精华

Rust TreeCore 独立 crate 用于 DynoSim 仿真 设计

ishandhanani 提议在仿真研究中无需 libtorch 依赖即可使用该 radix tree;dreamtalen 补充 DynoSim 是纯 Rust 离散事件仿真器,希望能直接复用本 PR 的 radix tree 以减少 SGLang 与其仿真之间的语义漂移。

结论:未在本 PR 落地,作为后续协作方向;PR 当前仍通过 PyO3 绑定 tch(libtorch)。 · 待处理

Python/Rust 后端等价性测试 测试

ispobock 建议:'The registry lets both run under one env var, so we could drive the same op sequence through both and assert equal device_indices, LRU eviction order, and CacheActions.' 同时提醒 rebase 到 main 并确认 Rust 逻辑与更新后的 Python 一致。

结论:以共享 UnifiedRadixCache 测试套件(2436 例)+ Python/Rust inspector 快照方式实现,而非后端各自复制测试。 · 已解决

重复 Full KV backup 问题 正确性

alphabetc1 指出 'A double-backup issue' 并附参考修复 PR,其余部分 LGTM。

结论:已通过独立 commit 'fix: prevent duplicate full kv backups'(作者 alphabetc1)修复。 · 已解决

Cargo 缓存目录并发竞态 other

ispobock 定位 ci_install_dependency.sh 共享 CARGO_TARGET_DIR,CUDA job 的 rm -rf 会打断同机 Rust 编译(deps/libuuid-*.rlib 被删),建议在 RUNNER_TEMP 下用 per-run target dir。

结论:最终 commit 'build rust exts in a per run target dir' 解决;此前也加了 'guard cargo target drop with a lock'。 · 已解决

风险与影响

  1. 语义漂移风险(高):Rust 后端是基于 Python 行为快照的移植,Python 侧后续新增语义(session、C128、external linker、custom component)尚未移植;共享测试套件覆盖 2436 例,但未覆盖这些组合,后续 Python 侧改动需同步移植并有回归风险。
  2. 部署兼容性风险(中):PR 明确 CUDA 13.4 / PyTorch 2.15 未启用,打包版本限定 PyTorch 2.11-2.13;AMD/MLX/NPU 生产部署未验证,只能 fail-fast。
  3. 构建链复杂性(中):tch 0.24 依赖 Torch 2.13 兼容头、cargo 并发构建、wheel repair 链路(prepare_sglang_wheel.py 新增 smoke test)都可能成为发布阻塞点。
  4. 回归面(中):buffer_mode/pipeline.py 从 UnifiedTreeNode 对象访问改为 BufferBackupSnapshot 快照 + validate_buffer_backup,影响 HiCache 备份管线;unified_tree_core.py 新增接口方法,Python 后端行为也有调整。
  5. 性能未验证:PR 动机包含性能,但未给出基准数据,Rust 后端的实际收益需后续 benchmark 确认。

对用户:默认路径仍是 Python 后端,Rust 后端通过环境变量/registry 选择,不支持的配置会显式报错而非静默错误,现有部署无破坏性变更。对系统:缓存核心新增一套无 Python 节点对象遍历的高性能后端候选,为 PD-decode HiCache、DFS 权重调度等热路径去除 Python 开销铺路。对团队:后续缓存语义迭代需同步维护双后端,共享测试套件降低回归风险但增加功能开发成本;CI 增加 Rust 构建、wheel smoke test 等环节。对外部社区:为 DynoSim 等 Rust 仿真生态提供了可复用 crate 的基础。

核心缓存路径双后端化 大规模语义移植需持续对齐 session/C128/linker 组合未覆盖 cargo 构建并发与缓存竞态 AMD/MLX/NPU 部署未验证 性能收益未附基准数据

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论