Prhub

#35644 [mem_cache][10/N] refactor: drop the redundant _component suffix in unified_cache/components

原始 PR 作者 alphabetc1 合并时间 2026-09-13 21:38 文件变更 18 提交数 8 评论 0 代码增减 +38 / -38

执行摘要

统一缓存组件去掉 _component 后缀并同步全部导入

动机来自 RFC #25371 的命名规则 §4.1:如果一个目录里的每个文件都共享目录已经声明的角色,这个词缀不携带信息,应以 models/qwen.py 而非 models/qwen_model.py 为范式。#32484 已经把 unified_cache_components/ 搬到 unified_cache/components/,却保留了 _component 后缀,于是目录名与文件名把同一个角色说了两遍,形成 unified_cache/components/full_component.py 这种冗余写法。PR body 还指出 tree_component.py 本身命名不准确:它不是「某个 tree component」,而是 TreeComponent ABC 加上所有组件都要 import 的共享枚举与工具函数,改叫 base.py 才能与 allocator/base.py、pool_host/base.py 的既有命名对齐。

如果只是想了解功能演进,可以跳过本 PR——它不改变任何运行时行为。但如果你在维护 unified_cache 相关代码、插件或下游集成,值得精读一次:重点是确认自己代码里的 import 路径,并留意「无兼容 shim」是该项目在 mem_cache 重构中的既定政策。值得关注的设计决策有两个:一是把 tree_component.py 改名为 base.py 而不是 tree.py,理由是它实际上是 ABC 与共享枚举的基座,这与 allocator/base.py、pool_host/base.py 的命名形成一致范式;二是用「字节级等价 + 前后测试计数零 delta」而非新增测试来验收机械搬迁,这套自证方法比「跑了 CI 且是绿的」更可审计,适合被其他重构 PR 借用。

讨论亮点

本 PR 的 review 记录非常轻:只有两条 APPROVED,没有一条 inline 评论,没有 review 线程。Jialin 的评价是 “Nice clean up!”,ispobock 空白批准。真正承载论证的是 PR body 自述的两条「自证」:一是命名口径(引用 #25371 §4.1,把 tree_component.py 改名 base.py 与 allocator/base.py、pool_host/base.py 对齐),二是等价性验收(字节级一致 + 前后测试计数零 delta)。因此这次合入的技术信号主要来自作者的自证材料,而不是 reviewer 的挑战;这也符合 RFC 规则 2「机械搬迁必须满足 mechanical-refactor-verify 技能」的既定流程。

实现拆解

  1. 文件搬迁(git mv,四个文件零符号改名):tree_component.py → base.py(承载 TreeComponent ABC、ComponentType、ComponentData、BASE_COMPONENT_TYPE、next_component_uuid 等共享定义),full_component.py → full.py(FullComponent),mamba_component.py → mamba.py(MambaComponent),swa_component.py → swa.py(SWAComponent)。用 git mv 是为了尽量保留 git blame,符合 RFC 规则 3。

  2. 包内导入重写:三个具体组件文件各自只改了一行 import 源,即 from ...components.tree_component import (...) → from ...components.base import (...);components/init.py 把 base 的再导出提到最前,三个具体组件按字母序跟在后面(isort 排序结果),all 内容一个字符未动,因此包级公共 API 保持不变。

  3. 包外调用点同步:cache_init_params.py、unified_cache/session_ref_tracker.py(TYPE_CHECKING 块内)、hardware_backend/mlx/kv_cache/auxiliary_state.py 等调用方全部换成新路径。auxiliary_state.py 的改动尤其值得注意——它是 MLX 硬件后端适配层,说明这次改名的影响面不止主仓 srt 路径。

  4. 跨语言引用同步:rust/sglang-radix-tree/src/components/mod.rs 中 TreeComponent trait 的 8 处 // Python reference — tree_component.py::TreeComponent.xxx 注释统一改为 base.py::...。只动注释,trait 方法与签名一行未变,Rust 侧二进制行为与接口完全不变。

  5. 测试与文档配套:8 个单测文件(test_mamba_path_state_cap.py、test_unified_cache_linker.py、test_tree_core_registry.py、test_mamba_donated_alloc_ratio.py、test_buffer_mode_sidecar.py、test_swa_locked_full_recover_unified.py、test_unified_radix_cache_bench.py、test_unified_radix_cache_unittest.py)只改 import;components/README.md 更新 5 行文件名引用;PR body 另称 test_session_unified_radix_cache.py 中 3 处按文件名断言文件内容的字符串字面量也已同步。

  6. 验收方式(无新增测试):作者没有新增用例,而是把验收标准改成字节级等价——base.py 与原 tree_component.py 完全一致,full.py / mamba.py / swa.py 各只差一行 import;并在 H200 上用 PYTHONPATH 遮蔽镜像内副本,对 5 个单测文件在 PR 分支与 pristine main 上跑出完全相同的 1117 passed / 1070 skipped / 52 subtests passed,据此证明 1070 个 skip 是既有基线而非本次引入。

文件 模块 状态 重要度
python/sglang/srt/mem_cache/unified_cache/components/base.py 统一缓存组件 renamed 5.15
python/sglang/srt/mem_cache/unified_cache/components/__init__.py 统一缓存组件 modified 5.47
rust/sglang-radix-tree/src/components/mod.rs 基数树内核 modified 5.32
python/sglang/srt/mem_cache/unified_cache/components/full.py 统一缓存组件 renamed 5.3
python/sglang/srt/hardware_backend/mlx/kv_cache/auxiliary_state.py MLX 后端 modified 4.87
test/registered/unit/mem_cache/test_mamba_path_state_cap.py 缓存单测 modified 4.2
python/sglang/srt/mem_cache/unified_cache/components/README.md 统一缓存组件 modified 2.37

关键符号

TreeComponent.component_type TreeComponent.create_match_validator TreeComponent.evict_component TreeComponent.acquire_component_lock TreeComponent.release_component_lock FullComponent.__init__ MambaComponent SWAComponent

关键源码片段

python/sglang/srt/mem_cache/unified_cache/components/__init__.py dependency-wiring

包级唯一再导出入口,是「不加 shim」策略下必须同步的公共面;本次除路径更新外还发生了 isort 重排(base 排到三个具体组件之前),是全部改动中唯一非纯路径的部分。

# 包级唯一再导出入口:改名后外部调用方仍然只依赖这些名字,
# 因此在「无兼容 shim」的约定下,这里是唯一必须同步的公共面。
from sglang.srt.mem_cache.unified_cache.components.base import (
    # 原 tree_component.py:TreeComponent ABC 与共享枚举、工具函数都从这里出。
    _NUM_COMPONENT_TYPES,
    BASE_COMPONENT_TYPE,
    CacheTransferPhase,
    ComponentData,
    ComponentType,
    EvictLayer,
    ExternalLinkerLoadPhase,
    LinkerTransferPhase,
    LRURefreshPhase,
    PrepareLoadBackResult,
    PreparePrefetchResult,
    TreeComponent,
    get_and_increase_time_counter,
    next_component_uuid,
)
# 三个具体组件按字母序排在 base 之后:isort 重排是本次唯一的非路径改动,
# 不改变任何运行时语义,但会让 diff 看起来比「纯改名」多一点噪音。
from sglang.srt.mem_cache.unified_cache.components.full import FullComponent
from sglang.srt.mem_cache.unified_cache.components.mamba import MambaComponent
from sglang.srt.mem_cache.unified_cache.components.swa import SWAComponent
​
__all__ = [
    "BASE_COMPONENT_TYPE",
    "ComponentData",
    "ComponentType",
    "ExternalLinkerLoadPhase",
    "LinkerTransferPhase",
    "EvictLayer",
    "FullComponent",
    "CacheTransferPhase",
    "LRURefreshPhase",
    "MambaComponent",
    "PrepareLoadBackResult",
    "PreparePrefetchResult",
    "SWAComponent",
    "TreeComponent",
    "_NUM_COMPONENT_TYPES",
    "next_component_uuid",
    "get_and_increase_time_counter",
]
rust/sglang-radix-tree/src/components/mod.rs documentation

TreeComponent trait 的 8 处 // Python reference 注释指向的 Python 文件名必须随改名同步,否则 Rust 侧读者会照着一个不存在的文件名去找 Python 实现;这是本 PR 唯一的跨语言改动面。

/// A per-component lock/value/eviction driver over the shared `UnifiedTreeCore`.
pub trait TreeComponent<K: ChildKeyType> {
    /// The component this driver serves.
    fn component_type(&self) -> ComponentType;
​
    /// Refresh this component's LRU position for `node_id` at the given walk phase.
    fn refresh_lru(
        &self,
        tree_core: &mut UnifiedTreeCore<K>,
        phase: LRURefreshPhase,
        node_id: NodeIdx_,
    ) {
        // Python reference — base.py::TreeComponent.refresh_lru:
        // 文件已由 tree_component.py 更名为 base.py,跨语言对照路径必须同步,
        // 否则 Rust 侧读者会照着不存在的文件名去 Python 里找实现。
        // match phase:
        // case LRURefreshPhase.WALKDOWN: ...
        // case LRURefreshPhase.MATCH_END: ...
        // case LRURefreshPhase.INSERT_END: return # WALKDOWN 已刷新整条插入路径
        unimplemented!("TreeComponent.refresh_lru")
    }
​
    /// Return a per-match stateful predicate deciding whether a node is a valid
    /// match boundary for this component.
    // Python reference — base.py::TreeComponent.create_match_validator:
    // 同一 trait 内共 8 处 reference 注释做了相同替换;trait 方法与签名一行未动,
    // 所以 Rust 侧的二进制接口与行为完全不变,纯注释同步。
    fn create_match_validator(
        &self,
        tree_core: &UnifiedTreeCore<K>,
        match_device_only: bool,
    ) -> Box<dyn FnMut(&UnifiedTreeCore<K>, NodeIdx_) -> bool>;
​
    /// Eviction priority on this node type; higher = evicted later, and evicting a
    /// component cascade-evicts every component of equal or lower priority.
    fn eviction_priority(&self, is_leaf: bool) -> i64 {
        0
    }
}
python/sglang/srt/mem_cache/unified_cache/components/full.py rename-or-move

full_component.py → full.py 的代表样本:整个文件只有一行 import 源变化,直观体现「除路径外逐字节等价」的搬迁承诺;FullComponent 的 component_type = ComponentType.FULL 等逻辑完全未动。

from __future__ import annotations
​
import heapq
from typing import TYPE_CHECKING, Callable, Optional, Sequence
​
import torch
​
from sglang.srt.mem_cache.base_prefix_cache import (
    DecLockRefParams,
    EvictParams,
    IncLockRefResult,
    InsertResult,
    MatchPrefixParams,
    MatchResult,
)
from sglang.srt.mem_cache.hicache_storage import (
    PoolName,
    PoolTransfer,
    PoolTransferResult,
)
from sglang.srt.mem_cache.unified_cache.cache_action import FreeComponentDeviceSlot
​
# 全文件唯一的实质改动:导入源从 components.tree_component 换成 components.base。
# 被导入的符号(TreeComponent、ComponentType、EvictLayer 等)一个都没改名,
# 所以 FullComponent 的实现与原 full_component.py 逐字节等价,只是模块路径变了。
from sglang.srt.mem_cache.unified_cache.components.base import (
    CacheTransferPhase,
    ComponentType,
    EvictLayer,
    ExternalLinkerLoadPhase,
    LinkerTransferPhase,
    TreeComponent,
)
​
​
class FullComponent(TreeComponent):
    component_type = ComponentType.FULL
​
    def __init__(self, cache, params):
        super().__init__(cache, params)
        # HiCache 状态:启用 HiCache 时该字段指向 host KV pool(主体逻辑本 PR 未触碰)。
        self._full_kv_pool_host = None

评论区精华

重命名是否符合 RFC 的命名规则(去掉不区分角色的词缀) 设计

作者在 PR body 引用 #25371 §4.1:如果目录内每个文件都共享目录已经声明的角色,该词缀不携带信息,应以 models/qwen.py 而非 models/qwen_model.py 为范式;并指出 tree_component.py 实际是 TreeComponent ABC 加共享枚举与工具,命名为 base.py 才能与 allocator/base.py、pool_host/base.py 对齐。Reviewer Jialin 回复 “Nice clean up!”。

结论:命名口径获得认可,四个文件的改名全部按 RFC 规则落地,未出现反对意见。 · 已解决

机械搬迁的等价性验收标准 测试

作者把验收标准从「功能正确」提升为「字节级等价」:base.py 与原 tree_component.py 完全一致,full.py / mamba.py / swa.py 各只差一行 import;并在 H200 上用 PYTHONPATH 遮蔽镜像内副本,对 5 个单测文件在 PR 分支与 pristine main 上跑出完全相同的 1117 passed / 1070 skipped / 52 subtests passed,据此证明 1070 个 skip 是既有基线而非本次引入。

结论:无需新增测试、也无需保留 re-export shim,机械搬迁通过验收。 · 已解决

跨语言引用同步:Rust trait 注释中的 Python reference documentation

rust/sglang-radix-tree/src/components/mod.rs 中 TreeComponent trait 的 8 处 // Python reference — tree_component.py::TreeComponent.xxx 全部改为 base.py::...;不改则 Rust 侧读者会照着已不存在的文件名去找 Python 实现,跨语言对照文档失真。改动仅限注释,trait 签名与实现零变化。

结论:注释路径同步完成,Rust 侧二进制行为与接口不变。 · 已解决

风险与影响

  1. 无兼容 shim 带来的导入面风险:按 RFC 规则 4,本 PR 不保留任何 re-export,因此任何未被覆盖的导入点会在 import 期直接 ImportError。RFC 明确把 hardware_backend/npu/、hardware_backend/mlx/ 列为 out of scope,而本 PR 只改到了 mlx 的 auxiliary_state.py——NPU 路径以及 benchmark 脚本、文档示例、下游插件中若还有 components.tree_component / components.mamba_component 之类的引用,会直接炸掉。

  2. 按文件名断言的测试存在文实不符的可能:PR body 声称改了 test_session_unified_radix_cache.py 中 3 处文件名字符串字面量(该测试按文件内容做断言),但在提供给我的 18 个变更文件清单里没有这个文件。若其断言确实存在而未被更新,测试会以「内容不匹配」的形式静默失败,而不是导入错误。这一点需要核实,属于上下文不足处。

  3. 目录布局仍可能回退:RFC §8 规划的 CI 守卫尚未落地,RFC 自己举的先例就是 allocator/ 目录存在两个月后仍有新文件落在顶层。本次改名在守卫缺位下完成,约束力靠人工评审维持。

  4. 分支长寿的维护成本:8 个 commit 中有 6 个是 merge main 解冲突(fix: resolve component rename conflicts with main 等),说明这条机械搬迁分支存活期间 main 反复改动同一批文件,diff 面积虽小但 rebase 成本被拉长,这类纯改名 PR 在冲突密集期本身是一种协作风险。

对最终用户:零影响,无运行时行为、性能、精度或显存占用变化,PR body 也明确 Benchmark 不适用。对核心开发者:所有 import sglang.srt.mem_cache.unified_cache.components.* 的代码与笔记本、调试脚本都需要跟着改,改动是纯文本替换但必须一次性完成,没有过渡期——这是 RFC 规则 4 刻意选择的代价,换来的是不允许新旧路径长期并存。对 Rust 侧维护者:TreeComponent trait 的 Python reference 注释恢复准确,跨语言对照文档不再指向一个已不存在的文件名;Rust 代码本身零变化。对团队流程:本 PR 演示了一套可复用的机械搬迁验收范式(git mv 保 blame + 字节级等价 + 同机同命令的前后测试计数对照),是 #25371 系列后续搬迁 PR 的模板。

核心模块路径变更 无兼容 shim(外部导入会直接失败) NPU 等 out-of-scope 路径待核对 测试按文件名做字符串断言 RFC §8 CI 守卫未落地

关联 Issue

#25371 [RFC][Refactor] `mem_cache` pool / allocator restructure

完整报告

参与讨论