Prhub

#49440 [Bugfix][KV Offload] Namespace persistent cache by model runner

原始 PR 作者 jongukc 合并时间 2026-07-26 13:21 文件变更 2 提交数 4 评论 6 代码增减 +22 / -1

执行摘要

为 KV 卸载持久缓存命名空间添加模型运行器标识

V1 和 V2 模型运行器的序列化缓存布局不兼容——V1 使用(head, layer, token, value)的跨层紧凑表示,而 V2 使用每层独立的张量(12 个张量,每个形状为 (blocks, heads, tokens, values))。尽管两种布局的字节数相同(均为 589824 bytes/row),但实际数据顺序不同,导致尺寸检查无法捕获不匹配,V2 加载 V1 缓存后会输出错误结果(例如从 Paris... 变成 .,,,,,, the the.,,,, by a)。该 PR 通过让非并行无关的配置文件路径不同来隔离 V1 和 V2 缓存。

建议详细阅读该 PR 的 issue 讨论和测试代码,了解 KV 卸载缓存布局的差异及处理方式。核心设计决策(将 parallel_agnostic 纳入哈希字段而非引入独立标识)是值得借鉴的模式。

讨论亮点

Reviewer orozery 最初建议不要在 FileMapper 中引入 model_runner 字段,而是将 parallel_agnostic 纳入哈希字段中,具体方式为在 fields 中添加 "parallel_agnostic": False。Author jongukc 采纳该建议并修改了实现。

实现拆解

  1. FileMapper.__init__ 中增加字段记录:当 parallel_agnosticFalse 时,在 self.fields 字典中添加 "parallel_agnostic": False(之前只有 parallel_agnostic=True 时才会对 fields 产生影响——将 tp/pp/pcp/dcp 置为 1 和 rank 置为 0)。由于 fields 用于计算持久化路径的哈希,这一变更使得非并行无关的配置产生与之前不同的路径哈希,从而自动与旧路径隔离。
  2. 调整测试用例断言:更新 test_get_run_config_fields 的预期输出,使其包含 "parallel_agnostic": False。更新 test_parallel_agnostic_collapses_namespace_when_config_allows 断言,确保 parallel_agnostic=True 且配置允许并行无关时,fields 中不包含 "parallel_agnostic" 键。更新 test_parallel_agnostic_ignored_when_config_disallowstest_namespace_kept_without_parallel_agnostic_opt_in,断言 fields["parallel_agnostic"]False
  3. 新增回归测试 test_parallel_agnostic_separates_persistent_layouts:创建一个并行无关的 mapper 和一个非并行无关的 mapper,验证它们的 base_path 不同,且 fields 中并行无关的 mapper 不包含 "parallel_agnostic" 键,而非并行无关的 mapper 包含 "parallel_agnostic": False
文件 模块 状态 重要度
vllm/v1/kv_offload/file_mapper.py 文件映射器 modified 5.69
tests/v1/kv_offload/test_file_mapper.py 测试 modified 5.71

关键符号

FileMapper.__init__ test_parallel_agnostic_separates_persistent_layouts

关键源码片段

vllm/v1/kv_offload/file_mapper.py core-logic

核心修复文件:在 `FileMapper.__init__` 中添加了条件判断 `if not parallel_agnostic: self.fields["parallel_agnostic"] = False`,使得非并行无关的配置在 fields 中包含该字段,从而改变持久化路径的哈希值,隔离 V1 和 V2 缓存。

# vllm/v1/kv_offload/file_mapper.pyclass FileMapper:
    def __init__(
        self,
        root_dir: str,
        model_name: str,
        tokens_per_hash: int,
        blocks_per_file: int,
        tp_size: int,
        pp_size: int,
        pcp_size: int,
        dcp_size: int,
        rank: int,
        dtype: str,
        kv_cache_groups: list[dict] | None = None,
        inference_engine: str = "vllm",
        parallel_agnostic: bool = False,
    ):
        """
        Initialize the file mapper. Each worker constructs its own, but
        `config.json` is shared across workers since rank lives outside the hash.
        When `parallel_agnostic=True`, tp/pp/pcp/dcp are forced to 1 and rank
        to 0 so multiple parallelism layouts collapse into the same folder.
        """
        # -- 并行无关时归一化并行参数到 1/0,避免路径差异 --
        if parallel_agnostic:
            tp_size = pp_size = pcp_size = dcp_size = 1
            rank = 0
        self.rank: int = rank
        # 核心变更:当非并行无关时,显式记录 `parallel_agnostic: False` 到 fields 中
        # 这会改变 `_compute_base_path` 的哈希结果,使 V2 的持久化路径
        # 自动与 V1 的旧路径隔离(V1 使用跨层紧凑布局,V2 使用逐层张量)
        self.fields: dict = {
            "model_name": model_name,
            "tokens_per_hash": tokens_per_hash,
            "blocks_per_file": blocks_per_file,
            "tp_size": tp_size,
            "pp_size": pp_size,
            "pcp_size": pcp_size,
            "dcp_size": dcp_size,
            "dtype": str(dtype),
            "kv_cache_groups": kv_cache_groups or [],
            "inference_engine": inference_engine,
        }
        # 新增行:只有显式声明非并行无关时,才加入该字段
        # 若为并行无关,该字段不出现,保持向后兼容
        if not parallel_agnostic:
            self.fields["parallel_agnostic"] = False
        self.base_path: str = self._compute_base_path(root_dir, self.fields)
tests/v1/kv_offload/test_file_mapper.py test-coverage

测试覆盖文件:更新已有测试断言以匹配新行为,并新增回归测试 `test_parallel_agnostic_separates_persistent_layouts`,验证并行无关和非并行无关的 mapper 产生不同 base_path。

# tests/v1/kv_offload/test_file_mapper.py ( 新增测试 )def test_parallel_agnostic_separates_persistent_layouts():
    # 创建并行无关的 mapper(模拟 V1:不使用 `parallel_agnostic: False`)
    agnostic = make_mapper_from_offloading_spec(
        is_parallelism_agnostic=True,
        parallel_agnostic=True,
    )
    # 创建非并行无关的 mapper(模拟 V2:会添加 `parallel_agnostic: False`)
    specific = make_mapper_from_offloading_spec(
        is_parallelism_agnostic=False,
        parallel_agnostic=True,
    )
​
    # 核心断言:两种布局应具有不同的 base_path
    assert agnostic.base_path != specific.base_path
    # 并行无关的 fields 中不应包含 `parallel_agnostic` 键
    assert "parallel_agnostic" not in agnostic.fields
    # 非并行无关的 fields 中应包含 `parallel_agnostic: False`
    assert specific.fields["parallel_agnostic"] is False

评论区精华

是否引入 model_runner 字段还是使用 `parallel_agnostic` 哈希 设计

Reviewer orozery 建议不要引入 `model_runner` 字段,而是将 `parallel_agnostic` 纳入 hashable fields 中,即在 fields 中添加 `"parallel_agnostic": False`。

结论:Author jongukc 采纳建议,修改实现为当 `not parallel_agnostic` 时在 fields 中添加 `"parallel_agnostic": False`。 · 已解决

风险与影响

低风险。核心逻辑仅增加一个条件判断(if not parallel_agnostic: self.fields["parallel_agnostic"] = False),不改变现有并行无关配置的行为。测试覆盖了主要路径和回归场景。潜在风险:任何依赖旧路径哈希的持久化缓存将在升级后失效,但这是预期行为——旧缓存本就是不安全的。

直接影响:修复了 V1/V2 模型运行器共享持久化 KV 卸载缓存时的静默数据损坏 bug。用户无需手动操作,升级后 V2 将自动使用隔离后的缓存路径。影响范围:所有使用 V1 KR 卸载引擎且持久化缓存已包含 V1 数据的用户,但 V2 运行器在升级后不会再误加载这些数据。

缓存路径变更 存量缓存失效

关联 Issue

#44733 [KV offload] Parallel-agnostic fs-tier cache for single full-attention group

完整报告

参与讨论