Prhub

#53624 [KV Offload] Add KVCR secondary-tier adapter

原始 PR 作者 mkhazraee 合并时间 2026-09-15 14:55 文件变更 7 提交数 14 评论 26 代码增减 +1211 / -1

执行摘要

新增 KVCR 二级卸载适配器并接入工厂与 CI

PR body 明确目标是「Add an in-tree adapter that exposes KVCR as a vLLM secondary tier」,并且划清边界:vLLM contains the factory registration and the translation between its secondary-tier interface and KVCR. KVCR remains an optional external package and owns the underlying storage, transfer, and policy implementation. Request-provided routing metadata is forwarded without making vLLM responsible for its semantics. 即把外部 KV 缓存系统的接入统一到已有的二级 tier 插件机制里,同时避免把 KVCR 的语义细节写进 vLLM。该 PR 堆叠在 #52067 与 #52068 之上,属于 KV offload 分层存储路线的延续。评审阶段又追加了两点约束:orozery 要求「change the unit tests to be skipped in case kvcr is not available」以及「Let's pin a specific version instead of @main」,前者是为了让可选依赖不阻塞 CI,后者是为了避免 mutable 分支带来的供应链风险(coderabbit 也以 CWE-494 提了同样的点)。

建议精读。对 KV offload / tiering 方向的维护者与打算编写自家二级 tier 适配器的团队,本 PR 是很好的模板:_FrameworkPinAdapter 展示了如何把外部系统的异步 pin 回调安全映射到主 tier 的请求生命周期(pending → process_pending → release_pin → take_pin_job_results),_VllmKeyAdapter 展示了 key 双向翻译时如何处理 group 维度与 int/bytes 两种事件格式,build_metric_definitions 展示了如何把外部指标并入 vLLM 命名空间。阅读时重点看三处设计决策:其一,注册惰性化 + 测试 importorskip 的可选依赖处理策略;其二,KVCR 语义(router hint、策略、描述符)全部外置的边界划分;其三,CI 固定 v0.1.0 tag 而非 @main 的供应链取舍。对只关心推理性能与模型支持的贡献者,可以略读,只需知道默认路径未变。另外建议 reviewer 单独盯一下 install-kv-offload.sh 里 uv pip uninstall 的容错,以及 _resolve_pin 静默吞异常的观测性。

讨论亮点
  1. 可选依赖与测试跳过:orozery 在 issue/PR 评论中要求「change the unit tests to be skipped in case kvcr is not available. This includes skipping kvcr in test_pre_registered_tiers_can_be_imported」,并在 test_factory.py:49 追加「Let's add a comment explaining this (kvcr is an optional external dependancy)」;作者以 commit「Skip unit tests in case KVCR is not available」「Clarifying KVCR is optional and the reason for the if」回应,最终形成 pytest.importorskip("kvcr") + find_spec("kvcr") 跳过 + 显式注释的组合。
  2. 软件供应链与版本固定:orozery 在 install-kv-connectors.sh 上评论「Let's pin a specific version instead of @main」,作者回复「we will tag the KVCR repo next week, for now without the tag went for main… will do this before merging」;coderabbit 也以 Security Misconfiguration(CWE-494) 建议固定到 commit。最终以 commit「Added KVCR repo tag」落到 v0.1.0 tag,并在最终 approve 里确认「LGTM Ready for merge once kvcr version is pinned」。
  3. primary_kv_view 是否需要 ndim 防御:coderabbit 在 manager.py:358 建议「Reject flat primary_kv_view inputs… Reject primary_kv_view.ndim < 2 before reading strides[0]」;作者反驳「The production path obtains this view from SharedOffloadRegion.create_kv_memoryview(), which explicitly constructs shape (num_chunks, row_stride_bytes). Filesystem, object-store, and example secondary tiers likewise rely on this invariant」,coderabbit 接受并记录学习项,线程以「不修改、沿用既有契约」收束。
  4. KVCR 语义该放在哪一侧:oandreeva-nv 在 manager.py 的导入块上评论「I think we need to hide it from vllm and move to kvcr」,指向 router hint 的解析逻辑;随后出现由她参与署名的 commit「Move KVCR router hint parsing out of vLLM」,最终 vLLM 只 import ROUTER_HINT_KEY / ROUTER_HINT_CAPABILITIES 并原样转发,语义解释权归还 KVCR。
  5. 演进中的外部 API:commit 历史里「Renamed block_ids to chunk_ids」「Update to multi-pool API (#4)」「Added nixl backend config support, and some rename updates」说明适配器是在 KVCR 尚未冻结的接口上迭代对齐的,评审过程本身承担了 API 对齐成本。

实现拆解

实现可拆成「注册 → 构造期配置翻译 → 运行期双向适配 → 测试 → CI 配套」五个环节:

  1. 注册入口(惰性加载):vllm/v1/kv_offload/tiering/factory.py 追加第 5 个内置 tier:SecondaryTierFactory.register_tier("kvcr", "vllm.v1.kv_offload.tiering.kvcr.manager", "KVCRSecondaryTierManager");get_tier_class 仍走 importlib.import_module + issubclass 校验,只有配置里出现 type: kvcr 才会真正 import 外部包。新增 vllm/v1/kv_offload/tiering/kvcr/__init__.py(仅许可证头)把适配器收进独立子包。这一步决定了未安装 KVCR 的部署完全不受影响,也决定了注册路径不会被默认 CI 自然覆盖。

  2. 构造期配置翻译与校验(vllm/v1/kv_offload/tiering/kvcr/manager.py 的 KVCRSecondaryTierManager.__init__):

    • 策略解析 _resolve_policy:内置 fifo / lru / g3_fifo / g3_lru 映射到 KVCR 的 FIFOPolicy / LRUPolicy / G3FIFOPolicy / G3LRUPolicy;其他名字必须含 .,经 resolve_obj_by_qualname 解析并校验是 KVCachePolicy 子类,否则抛 ValueError / TypeError。
    • 参数一致性:kvcr_service_socket_path 与 compatibility_digest 必须同时给出(异或直接 ValueError);router_capabilities 必须覆盖 _REQUIRED_ROUTER_CAPABILITIES(即 ROUTER_HINT_CAPABILITIES),否则 ValueError,保证请求级路由元数据能被 KVCR 正确消费。
    • KV events 一致性:events_enabled 与 self_describing_events、secondary_g2_slots、kvcr_service_socket_path、g3 的组合被显式校验,避免 vLLM 与 KVCR 各自产生不一致的块事件。
    • 拓扑对齐:控制面(control_host / control_ports / control_advertise_host / eager_ctrl_connect)按本地 DP rank 选择端口,KVCRGuardConfig 的 guard_index 同步为 data_parallel_rank_local,pool layout 改用主 tier 的 _primary_row_stride;local_dram_backend / remote_fw_dram_backend / g3 被翻译成 KVCR 的 KVCRBackendConfigs(G3 的字符串路径转 Path 并组装 G3Options)。
  3. 运行期双向适配:

    • 键格式:_VllmKeyAdapter.encode 只接受 bytes(否则 TypeError);decode 通过 get_offload_block_hash + maybe_convert_block_hash 把 vLLM 的 OffloadKey(含 group 维度)压回 block hash,并跟随 VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES 返回 int 或 bytes。
    • pin 生命周期:_FrameworkPinAdapter 把 KVCR 的 request_pin / process_pending / poll_pin_results / cancel_pin_request 回调翻译成主 tier 的 on_new_request → lookup(只保留 LookupResult.HIT)→ create_store_job → 逐 key 构造 MemDescriptor → on_request_finished 的完整请求生命周期;未命中的 key 以 None 占位,pin_handle 记为 kvcr-job-{job_id} 并登记到 _pin_jobs,release_pin 再换回成功的 JobResult。异常一律 logger.warning 后返回 None,不向调用方抛错。
    • 指标:build_metric_definitions 受 extra_config["enable_telemetry"] 控制,_kvcr_metric_definitions() 用 _vllm_metric_name(vllm: 前缀)把 KVCR 的 duration / bytes / blocks / state 指标映射为 vLLM 的 OffloadingHistogramMetadata / OffloadingCounterMetadata / OffloadingGaugeMetadata 并声明 labelnames。
  4. 测试配套:tests/v1/kv_offload/tiering/test_kvcr_tier.py 以 pytest.importorskip("kvcr") 开头,用 RecordingKVCR 桩记录 submit_hint / discard_hint / deliver / deposit / poll_completed / get_stats,并用 monkeypatch.setattr(kvcr_manager, "KVCR", ...) 与 ZmqPeerControlChannel 替换真实实现;覆盖 DP rank 端点选择、G3 路径转换、hint 原样转发与 discard_hint 清理、键解码的 group 折叠、descriptor 地址(_primary_base_addr + chunk_id * 16)、telemetry、policy 选择、部分完成与 drain。tests/v1/kv_offload/tiering/test_factory.py 的 CI sentinel test_pre_registered_tiers_can_be_imported 遍历 registry,对 tier_type == "kvcr" 且 find_spec("kvcr") is None 时 continue,并补了「KVCR is an optional external dependency」注释。

  5. CI 配套:新增 .buildkite/scripts/install-kv-offload.sh——CUDA_VERSION 为空时直接 exit 0;否则 uv pip install --system "kvcr @ git+https://github.com/ai-dynamo/kvcr.git@v0.1.0",再读取 uv pip show --system nixl 的版本号,卸载 nixl-cu12 / nixl-cu13 后按 nixl-cu${CUDA_VERSION%%.*} 重装同版本 wheel,保持与镜像 CUDA runtime 匹配。.buildkite/test_areas/misc.yaml 的「V1 KV Offload」作业在跑 v1/kv_offload、v1/simple_kv_offload 之前插入该脚本。评审早期版本曾把 KVCR 安装写在 install-kv-connectors.sh 并指向 @main,最终被迁移到独立脚本并固定到 v0.1.0 tag。

文件 模块 状态 重要度
vllm/v1/kv_offload/tiering/kvcr/manager.py 卸载适配层 added 9.08
tests/v1/kv_offload/tiering/test_kvcr_tier.py 卸载测试 added 7.91
vllm/v1/kv_offload/tiering/factory.py 二级存储层 modified 5.02
vllm/v1/kv_offload/tiering/kvcr/__init__.py 卸载适配层 added 4.12
tests/v1/kv_offload/tiering/test_factory.py 卸载测试 modified 4.48
.buildkite/scripts/install-kv-offload.sh 安装脚本 added 4.36
.buildkite/test_areas/misc.yaml 构建配置 modified 2.76

关键符号

KVCRSecondaryTierManager.__init__ KVCRSecondaryTierManager.build_metric_definitions _resolve_policy _vllm_metric_name _kvcr_metric_definitions _FrameworkPinAdapter.request_pin _FrameworkPinAdapter.process_pending _FrameworkPinAdapter.poll_pin_results _FrameworkPinAdapter._resolve_pin _FrameworkPinAdapter.release_pin _VllmKeyAdapter.encode _VllmKeyAdapter.decode SecondaryTierFactory.register_tier test_pre_registered_tiers_can_be_imported

关键源码片段

tests/v1/kv_offload/tiering/test_kvcr_tier.py test-coverage

544 行新增单测,用 RecordingKVCR 桩把适配器与真实传输解耦,覆盖配置翻译、hint 转发、键解码、descriptor 地址、telemetry、策略与 drain,是理解契约细节最直接的文档。

以下片段整合自 tests/v1/kv_offload/tiering/test_kvcr_tier.py,展示桩实现与一个端到端断言用例。

# kvcr 是可选外部依赖:整模块在未安装时直接 skip,而不是 fail
pytest.importorskip("kvcr")
​
​
class RecordingKVCR:
    """记录 KVCR 调用序列的桩,让适配器逻辑可以脱离真实传输层测试。"""
​
    def __init__(self) -> None:
        # ... 省略 config / guard_config / backend_configs 等记录字段
        self.submit_hint_calls: list[tuple[tuple[Any, ...], dict[str, Any]]] = []
        self.discard_hint_calls: list[str] = []
        self.deliver_calls: list[
            tuple[OpHandle, dict[BlockKey, list[MemDescriptor]], str | None]
        ] = []
        self.completed: list[tuple[OpHandle, dict[BlockKey, OpEntryResult]]] = []
        self._next_op_handle = 1
​
    def query(
        self,
        keys: Collection[BlockKey],
        request_id: str | None = None,
    ) -> list[tuple[QueryStatus, CacheTier | None]]:
        # 一个 query_status 开关即可驱动 HIT / FETCHING / FETCHABLE / MISS 四种返回
        tier = {
            QueryStatus.HIT: CacheTier.LOCAL_G2,
            QueryStatus.FETCHING: CacheTier.LOCAL_G2,
            QueryStatus.FETCHABLE: CacheTier.REMOTE_G2,
            QueryStatus.MISS: None,
        }[self.query_status]
        return [(self.query_status, tier) for _ in keys]
​
    def deliver(
        self,
        blocks: Mapping[BlockKey, list[MemDescriptor]],
        request_id: str | None = None,
    ) -> OpHandle:
        # 记录本次 load,并立刻把该 op 标记为全部 block 成功,模拟同步完成
        op_handle = self._next_op_handle
        self._next_op_handle += 1
        self.deliver_calls.append((op_handle, dict(blocks), request_id))
        self.complete(op_handle, {key: True for key in blocks})
        return op_handle
​
​
def test_kvcr_tier_adapts_request_and_load(monkeypatch):
    """Check hint forwarding, key conversion, load descriptors, and cleanup."""
    kvcr = RecordingKVCR()
    tier = _make_tier(monkeypatch, kvcr)
    router_hint = {"opaque": True}
    ctx = ReqContext(
        req_id="req",
        kv_transfer_params={ROUTER_HINT_KEY: router_hint, "unrelated": object()},
    )
    key = make_offload_key((123).to_bytes(8, "big"), 0)
    same_hash_other_group = make_offload_key((123).to_bytes(8, "big"), 7)
​
    tier.on_new_request(ctx)
​
    # 只转发 ROUTER_HINT_KEY 对应的那部分,vLLM 不解释 hint 的语义
    assert kvcr.submit_hint_calls == [((), {"request_id": "req", "hints": router_hint})]
​
    # decode 丢掉 group 维度:同 hash、不同 group 解码结果必须一致
    monkeypatch.setenv("VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES", "0")
    decode = kvcr.constructor_bindings.key_adapter.decode
    assert decode(BlockKey(bytes(key))) == (123).to_bytes(8, "big")
    assert decode(BlockKey(bytes(same_hash_other_group))) == (123).to_bytes(8, "big")
​
    tier.submit_load(_job(7, ctx, key=key, chunk_id=2))
​
    # descriptor 地址 = primary base + chunk_id * 行宽,保证 KVCR 直读正确的块
    (descriptor,) = kvcr.deliver_calls[0][1][key]
    assert descriptor.addr == tier._primary_base_addr + 2 * 16
    assert descriptor.end_point_name == kvcr.nixl_agent_name
    assert list(tier.get_finished_jobs()) == [JobResult(7, True)]
​
    # 请求结束必须显式 discard_hint,避免 KVCR 侧残留请求级元数据
    tier.on_request_finished(ctx)
    assert kvcr.discard_hint_calls == ["req"]
vllm/v1/kv_offload/tiering/factory.py core-logic

二级 tier 的唯一注册入口,新增 kvcr 一行注册决定了该 tier 能否被配置发现,也决定未安装 KVCR 时是否安全(惰性 import)。

# vllm/v1/kv_offload/tiering/factory.py
# kvcr 走惰性注册:这里只登记模块路径与类名,
# 只有配置真正选用 "kvcr" 时 get_tier_class 才会 import,
# 因此没安装 kvcr 的环境不会因为这条注册而启动失败。
SecondaryTierFactory.register_tier(
    "kvcr",
    "vllm.v1.kv_offload.tiering.kvcr.manager",
    "KVCRSecondaryTierManager",
)
tests/v1/kv_offload/tiering/test_factory.py test-coverage

CI sentinel 需要跳过未安装的 kvcr,否则可选依赖会让工厂注册测试在无 KVCR 环境里失败。

# tests/v1/kv_offload/tiering/test_factory.py
from importlib.util import find_spec
​
​
def test_pre_registered_tiers_can_be_imported():
    """CI sentinel: registered tiers import and yield SecondaryTierManager."""
    for tier_type in SecondaryTierFactory._registry:
        # KVCR is an optional external dependency and may not be installed.
        if tier_type == "kvcr" and find_spec("kvcr") is None:
            continue
        cls = SecondaryTierFactory._registry[tier_type]()
        assert issubclass(cls, SecondaryTierManager)

评论区精华

KVCR 语义(router hint 解析)应从 vLLM 移到 KVCR 设计

oandreeva-nv 在 `manager.py` 导入块评论「I think we need to hide it from vllm and move to kvcr」,指向 vLLM 侧解释 KVCR router hint 的做法;后续出现由她参与署名的 commit「Move KVCR router hint parsing out of vLLM」。

结论:router hint 解析被移出 vLLM,适配器只 import `ROUTER_HINT_KEY` / `ROUTER_HINT_CAPABILITIES` 并把请求携带的元数据原样转发,符合 PR body 中「forwarded without making vLLM responsible for its semantics」的定位。 · 已解决

KVCR 依赖版本固定(@main → 固定 tag/commit) 安全

orozery 在 `.buildkite/scripts/install-kv-connectors.sh` 评论「Let's pin a specific version instead of `@main`」;coderabbit 以 CWE-494(Download of Code Without Integrity Check)建议固定到具体 commit;作者回复「we will tag the KVCR repo next week… will do this before merging」。

结论:最终 commit「Added KVCR repo tag」把安装迁移到独立的 `install-kv-offload.sh` 并固定为 `kvcr @ git+…@v0.1.0`,orozery 在 approve 时确认「Ready for merge once kvcr version is pinned」。 · 已解决

可选依赖下的单元测试跳过 测试

orozery 要求「change the unit tests to be skipped in case kvcr is not available. This includes skipping kvcr in `test_pre_registered_tiers_can_be_imported`」,并在 `test_factory.py:49` 追加「Let's add a comment explaining this (kvcr is an optional external dependancy)」。

结论:`test_kvcr_tier.py` 用 `pytest.importorskip("kvcr")`,`test_factory.py` 的 sentinel 用 `find_spec("kvcr") is None` 跳过并补注释;对应 commit「Skip unit tests in case KVCR is not available」「Clarifying KVCR is optional and the reason for the if」。 · 已解决

primary_kv_view 是否需要 ndim 防御性校验 正确性

coderabbit 在 `manager.py:358` 建议「Reject `primary_kv_view.ndim < 2` before reading `strides[0]`」,理由是 flat view 会让 `_primary_row_stride` 与 `_make_descriptor` 算出错误地址;作者回应生产路径来自 `SharedOffloadRegion.create_kv_memoryview()`,形状固定为 `(num_chunks, row_stride_bytes)`,filesystem/obj/example tier 同样依赖该不变式。

结论:接受既有契约,不新增 KVCR 专属校验;coderabbit 确认「A KVCR-specific `ndim` guard is not needed」并记录为学习项,线程 resolve。 · 已决议(不修改)

风险与影响

  1. 外部包耦合与版本漂移:manager.py 对 KVCR 类型是硬 import(MemDescriptor、OpEntryResult、KVCRBindings、G3Options、KVCachePolicy 等),任何签名或语义变化都会直接在 KVCRSecondaryTierManager.__init__ 或 pin 回调路径上炸开。评审期内 KVCR 已经改过 block_ids → chunk_ids、KeyHintAdapter → KeyAdapter、多池 API 等;CI 固定 v0.1.0 tag 只能护住 CI,用户侧升级 KVCR 时缺少版本闸门(compatibility_digest 仅在配置了 service socket 的场景生效)。
  2. 失败被静默降级:_FrameworkPinAdapter._resolve_pin 用 except Exception 记 warning 后返回 None,on_request_finished 失败也只 warning,且这些失败没有对应的 telemetry 维度(_kvcr_metric_definitions 只有 duration / bytes / blocks / state)。生产中会表现为「KVCR 源数据拿不到」的黑洞,排障只能依赖日志。
  3. CI 脚本健壮性:install-kv-offload.sh 在 set -euo pipefail 下执行 uv pip uninstall --system nixl-cu12 nixl-cu13,没有 || true,只要两个包中有一个不存在就可能让该步骤失败;NIXL_VERSION=$(uv pip show --system nixl | sed -n 's/^Version: //p') 依赖 nixl 必然已安装且输出格式稳定;该脚本还会在同一作业里影响 v1/kv_offload、v1/simple_kv_offload 两个与 KVCR 无关的测试套件的运行环境(重装 nixl-cu* 系列)。
  4. 测试覆盖边界:单测用 RecordingKVCR 桩替换真实 KVCR,KVCR 侧的 buffer 注册、NIXL 传输、guard 协商、淘汰策略行为都不在 vLLM CI 覆盖内;pytest.importorskip 的语义是「装不上就 skip」,因此在 CUDA_VERSION 未设置(脚本提前 exit 0)的镜像上会静默少跑一整套 KVCR 测试。
  5. 契约依赖而非防御:_primary_row_stride / _make_descriptor 直接读取 primary_kv_view.strides[0],依赖 SharedOffloadRegion.create_kv_memoryview() 提供的 (num_chunks, row_stride_bytes) 二维形状。若未来出现绕过该构造的调用方,会算出错误的 DMA 地址与长度——概率低但后果严重,当前靠其他 tier 的同约定兜底。

对用户而言:需要把 KVCR 作为二级 tier 的部署现在有 in-tree 适配路径,policy / g3 / control_ports / local_dram_backend 等配置项可直接在 tier 配置里表达;不使用该 tier 的部署零影响(工厂惰性 import,未安装 kvcr 也不报错)。对系统而言:vllm/v1/kv_offload/tiering/kvcr/ 成为新的子系统边界,factory.py 的注册表被扩展;CI 侧「V1 KV Offload」作业新增一步依赖 Git tag 的网络安装,作业时长和外部可用性风险上升,且该步骤在 NVIDIA 与 AMD 镜像上都会执行(misc.yaml 的 mirror 继承 commands)。对团队与生态而言:确立了「外部 KV 存储项目 + vLLM 薄适配器」的协作模式,router hint 解析被移出 vLLM 之后,双方职责边界更清晰;代价是 vLLM 侧的契约测试无法独立于 KVCR 演进,后续 KVCR 大版本升级需要同步回归本 PR 的 544 行单测。

外部依赖耦合 静默降级失败路径 CI 脚本容错不足 缺少真实 KVCR 端到端测试 依赖二维 memoryview 契约

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论