执行摘要

• 一句话：修复 KV offload 事件元数据被过早丢弃
• 推荐动作：值得精读，属于小范围但典型的状态生命周期正确性修复。重点关注三点：record_store 中"先合并再覆盖"的防丢逻辑、_take_removed_event 由 pop 改为计数式 discard 的写法，以及 per-event removal_expected 标志相对 tier 属性的设计权衡（mkhazraee 对全局配置无法区分二级介质的解释）。对理解 KV offload 事件转发机制的工程师有直接参考价值。

功能与动机

PR body 明确说明："Track offloaded-block residencies by medium and ownership so removing one residency does not discard metadata still needed by another." 即 KV offload 场景下同一 block 可同时在 CPU 与二级存储驻留，而旧实现以 pop 方式在第一次 removal 就丢弃元数据，破坏了另一介质事件的自描述 payload（token_ids、block_hashes、ownership 等），进而影响依赖这些事件的 KV 消费者（如外部前缀缓存）。

实现拆解

本 PR 围绕 OffloadingEventsTracker 的元数据生命周期做小幅重构，核心思路是把"一条元数据"与"多个驻留"解耦，按最终驻留移除时才释放。

1. 扩展元数据模型：在 vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py 的 _OffloadEventMetadata 中新增 active_residencies: set[tuple[Medium, str | None]] 字段，用于记录每个介质 + 归属的存活驻留；_build_event_metadata 构造时初始化为 {(Medium.CPU, None)}，表示 CPU 主介质默认存活。
2. 修复覆盖丢失：record_store 与 _record_partial_tail 写入新快照前，先用 existing := self._pending_event_metadata.get(offload_key) 取出旧记录，通过 meta.active_residencies.update(existing.active_residencies) 合并旧驻留集合，避免直接把元数据重置回仅 CPU 的种子状态（这是 review 中 Claude 与 orozery 指出的 bug）。
3. 事件翻译改为计数式释放：_take_stored_event 在 event.removal_expected 为真时向 active_residencies 添加 (event.medium, event.ownership)；_take_removed_event 由 pop(key, None) 改为 get(key) + active_residencies.discard(...)，仅当集合为空时才真正 pop(key) 删除快照，保证另一个驻留仍可取用完整 payload。
4. 扩展事件接口：在 vllm/v1/kv_offload/base.py 的 OffloadingEvent 上新增可选字段 removal_expected: bool = False，由产生事件的介质方标记该 store 是否会伴随 removal；store-only 介质保持 False，避免 _pending_event_metadata 无限增长。该字段默认 False，向后兼容。
5. 测试配套：tests/v1/kv_connector/unit/offloading_connector/test_events.py 将原 test_secondary_stored_event_does_not_mutate_cpu_metadata 重构为参数化的 test_reoffload_preserves_secondary_residency，覆盖 record_store 与 record_partial_store 两条 re-offload 路径，并强化 test_pending_cpu_removal_consumes_hit_backfill_until_next_hit 验证 CPU removal 后 STORAGE 驻留的元数据仍保留；同时移除了对 ownership 的错误依赖断言。

关键文件：

• vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py（模块 事件追踪；类别 source；类型 core-logic；符号 _OffloadEventMetadata, OffloadingEventsTracker.record_store, OffloadingEventsTracker._record_partial_tail, OffloadingEventsTracker._build_event_metadata）: 核心修复文件：为 _OffloadEventMetadata 引入 active_residencies 集合，record_store / _record_partial_tail 合并既有驻留，_take_removed_event 由 pop 改为计数式 discard，实现元数据保留到最终驻留移除。
• tests/v1/kv_connector/unit/offloading_connector/test_events.py（模块 单元测试；类别 test；类型 test-coverage；符号 test_secondary_stored_event_does_not_mutate_cpu_metadata, test_reoffload_preserves_secondary_residency）: 测试配套：重构 test_secondary_stored_event_does_not_mutate_cpu_metadata 为参数化 test_reoffload_preserves_secondary_residency，覆盖 record_store 与 record_partial_store 两条 re-offload 路径，并强化 CPU removal 后元数据保留的断言。
• vllm/v1/kv_offload/base.py（模块 事件接口；类别 source；类型 core-logic；符号 OffloadingEvent）: 为 OffloadingEvent 新增 removal_expected 字段，作为 store-only 介质不发 removal 时的标记，防止 _pending_event_metadata 无限增长；是本次设计的接口基础。

关键符号：OffloadingEventsTracker.record_store, OffloadingEventsTracker.record_lookup, OffloadingEventsTracker.record_partial_store, OffloadingEventsTracker._record_partial_tail, OffloadingEventsTracker._build_event_metadata, OffloadingEventsTracker._take_stored_event, OffloadingEventsTracker._take_removed_event, OffloadingEventsTracker.take_events

关键源码片段

vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py

核心修复文件：为 _OffloadEventMetadata 引入 active_residencies 集合，record_store / _record_partial_tail 合并既有驻留，_take_removed_event 由 pop 改为计数式 discard，实现元数据保留到最终驻留移除。

class OffloadingEventsTracker:
    """从 store 到最终 eviction 追踪 offloaded chunks 的 KV 事件 payload。"""
​
    def record_store(
        self,
        req: Request,
        group_config: "GroupOffloadConfig",
        chunk_idx: int,
        offload_key: OffloadKey,
    ) -> None:
        """为某个 offloaded chunk 快照 KV 事件 payload。"""
        if not self.self_describing_enabled:
            return
        if group_config.sliding_window_size_in_chunks is not None:
            return
        meta = self._build_event_metadata(req, group_config, chunk_idx)
        # 关键修复：直接覆盖会把 active_residencies 重置回仅 CPU 的种子，
        # 令已存在的 STORAGE 驻留被遗忘；这里先合并旧集合再写入。
        if existing := self._pending_event_metadata.get(offload_key):
            meta.active_residencies.update(existing.active_residencies)
        self._pending_event_metadata[offload_key] = meta
​
    def _take_removed_event(self, event: OffloadingEvent) -> Iterable[KVCacheEvent]:
        """把 removal 事件翻译为 BlockRemoved，并在最后一次驻留移除后清理。"""
        locality = event.locality.value if event.locality is not None else None
        by_group: dict[int, list] = {}
        for key in event.keys:
            # 由 pop 改为 get：一个驻留的移除不应丢弃其他介质仍需要的元数据
            meta = self._pending_event_metadata.get(key)
            if meta is not None:
                group_idx = meta.group_idx
                by_group.setdefault(group_idx, []).extend(
                    maybe_convert_block_hash(h) for h in meta.block_hashes
                )
                # 只移除本事件对应的 (medium, ownership) 驻留记录
                meta.active_residencies.discard((event.medium, event.ownership))
                if not meta.active_residencies:
                    # 所有驻留都已移除，才真正删除快照
                    self._pending_event_metadata.pop(key)
            else:
                # 未跟踪 key：保留原有告警 / 占位行为（本 PR 未改动）
                if self.self_describing_enabled:
                    logger.warning_once(...)
        # 后续按 group 构造 BlockRemoved 事件（此处从略）

tests/v1/kv_connector/unit/offloading_connector/test_events.py

测试配套：重构 test_secondary_stored_event_does_not_mutate_cpu_metadata 为参数化 test_reoffload_preserves_secondary_residency，覆盖 record_store 与 record_partial_store 两条 re-offload 路径，并强化 CPU removal 后元数据保留的断言。

@pytest.mark.parametrize(
    ("record_method", "position"),
    [("record_store", 0), ("record_partial_store", 4)],
)
def test_reoffload_preserves_secondary_residency(record_method, position):
    # 先让同一 key 以 STORAGE 介质 + custom 归属驻留（removal_expected=True）
    tracker, req, group_config, key = _lookup_chunk()
    [stored] = tracker.take_events(
        [
            _stored_event(
                [key],
                Medium.STORAGE,
                ownership="custom",
                removal_expected=True,
            )
        ]
    )
    # 再次 offload 到 CPU（record_store 或 record_partial_store），
    # 验证新的元数据覆盖不会丢掉已有的 STORAGE 驻留
    getattr(tracker, record_method)(req, group_config, position, key)
    assert tracker._pending_event_metadata[key].active_residencies == {
        (Medium.CPU, None),
        (Medium.STORAGE, "custom"),
    }
    # STORAGE 的 removal 只应移除它自己的驻留，
    # 元数据仍需保留给 CPU 驻留后续使用
    [removed] = tracker.take_events(
        [_removed_event([key], Medium.STORAGE, ownership="custom")]
    )
    assert stored.token_ids == [1, 2, 3, 4]
    assert stored.ownership == removed.ownership == "custom"
    assert key in tracker._pending_event_metadata

vllm/v1/kv_offload/base.py

为 OffloadingEvent 新增 removal_expected 字段，作为 store-only 介质不发 removal 时的标记，防止 _pending_event_metadata 无限增长；是本次设计的接口基础。

@dataclass
class OffloadingEvent:
    keys: list[OffloadKey]
    medium: Medium
    # True if blocks are removed, False if stored
    removed: bool
    locality: Locality | None = None
    # 次级介质标识；None 表示 primary（CPU）
    ownership: str | None = None
    # 新增字段：该 store 事件是否会在未来伴随 removal 事件。
    # store-only 介质（如 FS、对象存储）不会发出 BlockRemoved，
    # 保持 False 可避免驻留计数导致元数据无限增长。
    removal_expected: bool = False

评论区精华

核心讨论围绕两个正确性/设计问题，均由 orozery 提出：

1. record_store 覆盖导致驻留丢失：orozery 引用 Claude 的分析指出："record_store overwrites the whole metadata entry, which resets active_residencies back to the CPU-only seed. If a key already tracked as STORAGE-resident is re-offloaded to CPU, the STORAGE residency is forgotten." mkhazraee 确认并修复（commit 8cdac81），同时补充了针对性测试。
2. store-only 介质导致元数据无限增长：orozery 指出："Some tiers are store only (e.g. FS and OBJ). i.e. they don't emit BlockRemoved. This will cause the _pending_event_metadata dict to grow indefinitely." mkhazraee 解释设计取舍：将能力放回 SecondaryTierManager 需要新增大量 manager-to-tracker 管道代码，且 OffloadingKVEventsConfig 是全局的无法区分二级介质，因此选择在 OffloadingEvent 上新增 removal_expected 字段。orozery 接受并建议后续跟进："Adding per-event flag is fine for now. We should follow-up in making it a tier property."

• record_store 覆盖 active_residencies 导致 STORAGE 驻留丢失 (correctness): mkhazraee 确认问题并推送修复（commit 8cdac81）：record_store 与 _record_partial_tail 在覆盖前合并既有 active_residencies，并补充 test_reoffload_preserves_secondary_residency 覆盖两条路径。
• store-only 介质不产生 removal 导致元数据无限增长 (design): orozery 接受临时方案并建议 follow-up："Adding per-event flag is fine for now. We should follow-up in making it a tier property."

风险与影响

• 风险：主要风险集中在事件生命周期语义变化与增长隐患：
  • 依赖 removal_expected 正确设置：active_residencies 的清理完全依赖事件发起方正确设置 removal_expected。若某介质实际不产生 removal 却误标 True，元数据仍可能无限增长；当前靠调用方纪律缓解，orozery 已明确要求后续将"是否移除"提升为 tier 属性。
  • 事件流语义变更：CPU removal 不再立即清空元数据，外部消费者若依赖"removal 即清理"的旧行为，会观察到 _pending_event_metadata 生命周期变长；事件顺序本身不变，但同 key 的后续事件将携带完整 payload 而非 placeholder，行为差异可能影响依赖占位语义的对接方。
  • 测试覆盖局限：本次仅新增单元测试，未覆盖真实多介质（如 FS/对象存储 + CPU）的端到端事件流，removal_expected 的传播正确性依赖各介质调用点，建议后续集成验证。
  • 影响：影响范围限定在启用 KV cache events 且开启 self-describing events 的 offload 部署：修复了多介质驻留下 BlockStored / BlockRemoved 事件 payload（token_ids、block_hashes、ownership）丢失或退化为占位符的问题，对仅 CPU 单驻留场景行为不变（discard 后集合为空即 pop）。对团队而言，这是 KV offload 事件模型向"按 (medium, ownership) 管理驻留生命周期"演进的一步，为后续把 removal 预期收敛为 tier 属性打下基础；API 层面仅新增默认 False 的可选字段，向后兼容。
  • 风险标记：多介质驻留语义变更, 依赖 removal_expected 正确设置, 缺失端到端集成测试

关联脉络

• PR #52067 [KV Offload] 前置事件翻译基础（PR body 依赖声明）: PR body 明确声明 Depends on PR #52067，本 PR 在其 offloading 事件翻译能力之上扩展多介质驻留语义。
• PR #54284 [Bugfix][V1] Keep an encoder cache entry until its last occurrence is freed: 同一类"元数据保留到最后一个引用释放"的生命周期正确性修复，印证 v1 缓存层正在系统性地从一次释放改为引用计数式释放。
• PR #51358 [Bugfix][Mooncake] Save exact Mamba boundary states: 同属 offloading 事件与前缀缓存正确性线，修复事件存储错误导致外部前缀缓存中毒，与本 PR 共享 offloading scheduler / event 语义上下文。