#51161 [Bugfix][KV Offload] Handle chunked local attention in offloading scheduler

原始 PR 作者 almogtavor 合并时间 2026-08-09 15:25 文件变更 2 提交数 7 评论 13 代码增减 +24 / -0

执行摘要

KV offload 支持 chunked local attention，修复 Llama 4 启动崩溃

关联 Issue #51150 明确指出：任何使用 chunked local attention（即 Llama 4 家族）的模型在启用 offloading connector 时都会在启动阶段崩溃，原因就是 get_sliding_window_size_in_blocks() 只处理 SlidingWindowSpec 和 MambaSpec，然后断言其余均为 FullAttentionSpec。PR body 给出了关键等价性论证：chunked local attention 永远不会回看超过一个 attention chunk，因此其可达尾部与大小为 attention_chunk_size 的滑动窗口完全相同；并且强调向下取整会"静默跳过本可服务命中的 chunk"，所以必须向上取整。

值得精读。本 PR 是理解 vLLM KV offload 调度器如何为不同注意力架构估算"可达尾部"的理想入口：重点看 get_sliding_window_size_in_chunks() 与 is_store_reachable_swa_chunk() 的配合，以及 cdiv 向上取整背后"宁可多存、不可漏存"的安全方向选择。改动小、审查聚焦，适合作为团队内部学习 KV offload 调度机制的样例。

讨论亮点

核心讨论围绕三点：（1）orozery 在 issue 中追问 offloading 在 Llama 4 上是否真的验证过（"Did we verify offloading works with llama4 after this fix?"），almogtavor 随后用 Llama 4 Scout FP8、TP=4、H100 给出端到端数据：存储与读回各 245 MB、ext_hits 1248/13340、输出一致，并强调 eviction 由无关流量驱动而非 reset_prefix_cache（后者会连 connector 索引一起丢弃）。（2）Etelis 首轮 CHANGES_REQUESTED，要求删除冗余注释（"Remove useless comments"）并质疑兜底断言为何要展开（"Why do we need to elaborate here?"），最终代码保留单行注释与简洁断言。（3）Etelis 认为 "One test is enough IMO"，但作者保留两个断言以分别覆盖整除与部分 chunk 向上取整场景，reviewer 最终 APPROVED。

实现拆解

核心逻辑变更（vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py）：在 get_sliding_window_size_in_chunks() 的 SlidingWindowSpec 分支之后、MambaSpec 分支之前插入 ChunkedLocalAttentionSpec 分支——先断言 attention_chunk_size > 0 兜底，再返回 cdiv(attention_chunk_size, tokens_per_chunk)。cdiv 向上取整是关键：若向下取整会低估可达尾部，调度器会静默跳过仍可能命中外部缓存的 chunk。
兜底断言保留：函数末尾仍保留对 FullAttentionSpec 的断言，防止未来新增注意力 spec 时被静默当作全注意力处理（评审过程中曾扩写为带消息断言，后按 reviewer 意见精简）。
测试配套（tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py）：新增 test_chunked_local_attention_reports_its_chunk_window()，构造 attention_chunk_size=8192 的 ChunkedLocalAttentionSpec，验证 tokens_per_chunk=1024 时返回 8（整除场景）、tokens_per_chunk=3000 时返回 3（2.73 向上取整），覆盖取整边界。
端到端验证（PR body 记录，非代码改动）：先用 issue 中的 weight-free Llama 4 配置在 H100 单卡复现"stock 崩溃 / patched 启动成功"，再用 RedHatAI/Llama-4-Scout-17B-16E-Instruct-FP8-dynamic、TP=4 验证 KV 真实往返：245 MB 读回量与存储量一致，输出逐字节一致；还用小池 + 无关流量驱动 eviction 验证了完整 offload 生命周期。

文件	模块	状态	重要度
`vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py`	KV 卸载	modified	5.8
`tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py`	调度器	modified	5.07

关键符号

get_sliding_window_size_in_chunks test_chunked_local_attention_reports_its_chunk_window

关键源码片段

vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py core-logic

核心修复文件。get_sliding_window_size_in_chunks() 新增 ChunkedLocalAttentionSpec 分支，将 Llama 4 的可达尾部等价为 attention_chunk_size 大小的滑动窗口并向上取整，是该 PR 唯一的源码逻辑变更。

# 返回当前注意力组 " 可回看 " 的尾部大小（以 chunk 为单位）。
# None 表示全注意力（FullAttentionSpec），尾部不受限。
def get_sliding_window_size_in_chunks(
    kv_cache_spec: KVCacheSpec, tokens_per_chunk: int
) -> int | None:
    if isinstance(kv_cache_spec, SlidingWindowSpec):
        assert kv_cache_spec.sliding_window > 0
        return cdiv(kv_cache_spec.sliding_window, tokens_per_chunk)

    if isinstance(kv_cache_spec, ChunkedLocalAttentionSpec):
        # ChunkedLocalAttentionSpec（Llama 4 家族）注意力不会回看超过一个
        # attention chunk，因此可达尾部等价于 attention_chunk_size 大小的
        # 滑动窗口。cdiv 向上取整保证尾部永不被低估——低估会漏掉本可命中
        # 外部缓存的 chunk，造成静默缓存缺失。
        assert kv_cache_spec.attention_chunk_size > 0
        return cdiv(kv_cache_spec.attention_chunk_size, tokens_per_chunk)

    if isinstance(kv_cache_spec, MambaSpec):
        # Mamba 只依赖单个状态，尾部计为 1 个 chunk
        return 1

    # 兜底：保留对 FullAttentionSpec 的显式校验，防止未来新增注意力架构时
    # 被静默当作全注意力处理
    assert isinstance(kv_cache_spec, FullAttentionSpec)
    return None

tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py test-coverage

新增回归测试，锁定 ChunkedLocalAttentionSpec 分支的返回值，覆盖整除与部分 chunk 向上取整两种场景，是防止该修复回退的关键配套。

def test_chunked_local_attention_reports_its_chunk_window():
    """Llama 4 使用 chunked local attention，之前会触发 FullAttentionSpec
    断言导致引擎启动崩溃；本测试锁定修复后的返回值。"""
    spec = ChunkedLocalAttentionSpec(
        block_size=16,
        num_kv_heads=2,
        head_size=64,
        dtype=torch.bfloat16,
        attention_chunk_size=8192,
    )

    # 整除场景：8192 / 1024 = 8 个 chunk
    assert get_sliding_window_size_in_chunks(spec, tokens_per_chunk=1024) == 8
    # 部分 chunk 场景：8192 / 3000 ≈ 2.73，cdiv 向上取整为 3，
    # 保证可达尾部不会因向下取整而被低估
    assert get_sliding_window_size_in_chunks(spec, tokens_per_chunk=3000) == 3

评论区精华

Llama 4 上 offloading 是否端到端验证 question

orozery 在 issue 中追问 "Did we verify offloading works with llama4 after this fix?"，almogtavor 随后用 RedHatAI/Llama-4-Scout-17B-16E-Instruct-FP8-dynamic、TP=4、H100 验证：245 MB 读回量与存储量匹配，ext_hits 1248/13340，输出逐字节一致；并说明 eviction 由无关流量驱动而非 reset_prefix_cache（后者会连 connector 索引一起丢弃）。

结论：端到端验证通过，含真实权重与逐字节一致性确认。 · 已解决

新增分支的注释冗余 style

Etelis 在 review 中要求 "Remove useless comments"，针对 ChunkedLocalAttentionSpec 分支的多行解释性注释。

结论：最终代码保留单行简洁注释 # Attention never reaches back past one chunk。 · 已解决

兜底断言是否需展开描述 style

Etelis 针对把裸 assert 扩写为带消息断言的改动提问 "Why do we need to elaborate here?"。

结论：按 reviewer 意见精简，最终兜底断言保持简洁，未保留冗长消息。 · 已解决

测试断言数量 测试

Etelis 评论 "One test is enough IMO"，针对测试函数中的两个断言。

结论：作者保留两个断言，分别覆盖整除与部分 chunk 向上取整两个场景，reviewer 最终 APPROVED。 · 已解决

风险与影响

回归面小但链路关键：改动只影响 ChunkedLocalAttentionSpec 分支，SlidingWindowSpec/MambaSpec/FullAttentionSpec 路径行为不变；但 get_sliding_window_size_in_chunks() 的返回值会被 is_store_reachable_swa_chunk() 消费（后者有 assert sliding_window_chunks is not None），返回值错误会直接导致调度断言失败或缓存命中率异常。
取整方向安全：cdiv 向上取整只会高估可达尾部，代价是多存少量不可命中的 chunk（存储开销），而不会漏存，避免静默缓存缺失。
测试覆盖局限：单元测试只验证辅助函数返回值，未在 CI 中覆盖 Llama 4 + offload 的组合；端到端验证依赖手工 H100 实验，且只验证了 Scout，Behemoth 与更长上下文场景未覆盖。
非 offload 路径无影响：该函数仅在 offloading connector 调度器内被调用。

用户侧：Llama 4 家族（Scout/Behemoth）此前开启 KV offload 必然启动崩溃，本 PR 直接解锁该场景，用户可借助 CPU offload 扩展 KV 容量。系统侧：KV offload 的注意力 spec 覆盖从 sliding-window/mamba/full-attention 扩展到 chunked local attention，补上了 kv-connector 功能线的一块硬缺口。团队侧：提供了一个可复制的扩展模式——在 spec 分类函数中为新的注意力架构建模并保留兜底断言，后续接入新架构时可参照。影响程度中等：修复的是启动崩溃与功能解锁，代码面窄，但受益模型面广。

核心调度逻辑变更缺少 Llama 4 + offload 组合 CI 测试端到端验证依赖手工 H100 实验

关联 Issue

#51150 [Bug]: OffloadingConnector crashes at startup for Llama 4 models (ChunkedLocalAttentionSpec unhandled)

完整报告

执行摘要

一句话：KV offload 支持 chunked local attention，修复 Llama 4 启动崩溃
推荐动作：值得精读。本 PR 是理解 vLLM KV offload 调度器如何为不同注意力架构估算"可达尾部"的理想入口：重点看 get_sliding_window_size_in_chunks() 与 is_store_reachable_swa_chunk() 的配合，以及 cdiv 向上取整背后"宁可多存、不可漏存"的安全方向选择。改动小、审查聚焦，适合作为团队内部学习 KV offload 调度机制的样例。

功能与动机

实现拆解

核心逻辑变更（vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py）：在 get_sliding_window_size_in_chunks() 的 SlidingWindowSpec 分支之后、MambaSpec 分支之前插入 ChunkedLocalAttentionSpec 分支——先断言 attention_chunk_size > 0 兜底，再返回 cdiv(attention_chunk_size, tokens_per_chunk)。cdiv 向上取整是关键：若向下取整会低估可达尾部，调度器会静默跳过仍可能命中外部缓存的 chunk。
兜底断言保留：函数末尾仍保留对 FullAttentionSpec 的断言，防止未来新增注意力 spec 时被静默当作全注意力处理（评审过程中曾扩写为带消息断言，后按 reviewer 意见精简）。
测试配套（tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py）：新增 test_chunked_local_attention_reports_its_chunk_window()，构造 attention_chunk_size=8192 的 ChunkedLocalAttentionSpec，验证 tokens_per_chunk=1024 时返回 8（整除场景）、tokens_per_chunk=3000 时返回 3（2.73 向上取整），覆盖取整边界。
端到端验证（PR body 记录，非代码改动）：先用 issue 中的 weight-free Llama 4 配置在 H100 单卡复现"stock 崩溃 / patched 启动成功"，再用 RedHatAI/Llama-4-Scout-17B-16E-Instruct-FP8-dynamic、TP=4 验证 KV 真实往返：245 MB 读回量与存储量一致，输出逐字节一致；还用小池 + 无关流量驱动 eviction 验证了完整 offload 生命周期。

关键文件：

vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py（模块 KV卸载；类别 source；类型 core-logic；符号 get_sliding_window_size_in_chunks）: 核心修复文件。get_sliding_window_size_in_chunks() 新增 ChunkedLocalAttentionSpec 分支，将 Llama 4 的可达尾部等价为 attention_chunk_size 大小的滑动窗口并向上取整，是该 PR 唯一的源码逻辑变更。
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py（模块调度器；类别 test；类型 test-coverage；符号 test_chunked_local_attention_reports_its_chunk_window）: 新增回归测试，锁定 ChunkedLocalAttentionSpec 分支的返回值，覆盖整除与部分 chunk 向上取整两种场景，是防止该修复回退的关键配套。

关键符号：get_sliding_window_size_in_chunks, test_chunked_local_attention_reports_its_chunk_window

关键源码片段

`vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py`

# 返回当前注意力组 " 可回看 " 的尾部大小（以 chunk 为单位）。
# None 表示全注意力（FullAttentionSpec），尾部不受限。
def get_sliding_window_size_in_chunks(
    kv_cache_spec: KVCacheSpec, tokens_per_chunk: int
) -> int | None:
    if isinstance(kv_cache_spec, SlidingWindowSpec):
        assert kv_cache_spec.sliding_window > 0
        return cdiv(kv_cache_spec.sliding_window, tokens_per_chunk)

    if isinstance(kv_cache_spec, ChunkedLocalAttentionSpec):
        # ChunkedLocalAttentionSpec（Llama 4 家族）注意力不会回看超过一个
        # attention chunk，因此可达尾部等价于 attention_chunk_size 大小的
        # 滑动窗口。cdiv 向上取整保证尾部永不被低估——低估会漏掉本可命中
        # 外部缓存的 chunk，造成静默缓存缺失。
        assert kv_cache_spec.attention_chunk_size > 0
        return cdiv(kv_cache_spec.attention_chunk_size, tokens_per_chunk)

    if isinstance(kv_cache_spec, MambaSpec):
        # Mamba 只依赖单个状态，尾部计为 1 个 chunk
        return 1

    # 兜底：保留对 FullAttentionSpec 的显式校验，防止未来新增注意力架构时
    # 被静默当作全注意力处理
    assert isinstance(kv_cache_spec, FullAttentionSpec)
    return None

`tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py`

新增回归测试，锁定 ChunkedLocalAttentionSpec 分支的返回值，覆盖整除与部分 chunk 向上取整两种场景，是防止该修复回退的关键配套。

def test_chunked_local_attention_reports_its_chunk_window():
    """Llama 4 使用 chunked local attention，之前会触发 FullAttentionSpec
    断言导致引擎启动崩溃；本测试锁定修复后的返回值。"""
    spec = ChunkedLocalAttentionSpec(
        block_size=16,
        num_kv_heads=2,
        head_size=64,
        dtype=torch.bfloat16,
        attention_chunk_size=8192,
    )

    # 整除场景：8192 / 1024 = 8 个 chunk
    assert get_sliding_window_size_in_chunks(spec, tokens_per_chunk=1024) == 8
    # 部分 chunk 场景：8192 / 3000 ≈ 2.73，cdiv 向上取整为 3，
    # 保证可达尾部不会因向下取整而被低估
    assert get_sliding_window_size_in_chunks(spec, tokens_per_chunk=3000) == 3

评论区精华

Llama 4 上 offloading 是否端到端验证 (question): 端到端验证通过，含真实权重与逐字节一致性确认。
新增分支的注释冗余 (style): 最终代码保留单行简洁注释 # Attention never reaches back past one chunk。
兜底断言是否需展开描述 (style): 按 reviewer 意见精简，最终兜底断言保持简洁，未保留冗长消息。
测试断言数量 (testing): 作者保留两个断言，分别覆盖整除与部分 chunk 向上取整两个场景，reviewer 最终 APPROVED。

风险与影响

风险：
1. 回归面小但链路关键：改动只影响 ChunkedLocalAttentionSpec 分支，SlidingWindowSpec/MambaSpec/FullAttentionSpec 路径行为不变；但 get_sliding_window_size_in_chunks() 的返回值会被 is_store_reachable_swa_chunk() 消费（后者有 assert sliding_window_chunks is not None），返回值错误会直接导致调度断言失败或缓存命中率异常。
2. 取整方向安全：cdiv 向上取整只会高估可达尾部，代价是多存少量不可命中的 chunk（存储开销），而不会漏存，避免静默缓存缺失。
3. 测试覆盖局限：单元测试只验证辅助函数返回值，未在 CI 中覆盖 Llama 4 + offload 的组合；端到端验证依赖手工 H100 实验，且只验证了 Scout，Behemoth 与更长上下文场景未覆盖。
4. 非 offload 路径无影响：该函数仅在 offloading connector 调度器内被调用。
  - 影响：用户侧：Llama 4 家族（Scout/Behemoth）此前开启 KV offload 必然启动崩溃，本 PR 直接解锁该场景，用户可借助 CPU offload 扩展 KV 容量。系统侧：KV offload 的注意力 spec 覆盖从 sliding-window/mamba/full-attention 扩展到 chunked local attention，补上了 kv-connector 功能线的一块硬缺口。团队侧：提供了一个可复制的扩展模式——在 spec 分类函数中为新的注意力架构建模并保留兜底断言，后续接入新架构时可参照。影响程度中等：修复的是启动崩溃与功能解锁，代码面窄，但受益模型面广。
  - 风险标记：核心调度逻辑变更, 缺少 Llama 4 + offload 组合 CI 测试, 端到端验证依赖手工 H100 实验

关联脉络

PR #48414 [KV Connector] Canonical CPU layout for parallelism-agnostic KV offload: KV offload 功能线的架构基准 PR，本 PR 修复的正是该架构中调度器侧对注意力 spec 分类的缺口。
PR #51213 [XPU][Test] Pin block size in test_multi_connector: 同为 kv_connector 单元测试区域的跨平台修复，说明该测试基座在持续演进。
PR #47030 [ROCm][DistInf] Enable vLLM DI CI with buildkite/slurm: 分体式推理（disagg）CI 基建，与 offloading connector 同属 KV 传输/卸载领域，体现该领域的 CI 覆盖扩展方向。

#51161 [Bugfix][KV Offload] Handle chunked local attention in offloading scheduler

执行摘要

KV offload 支持 chunked local attention，修复 Llama 4 启动崩溃

实现拆解

评论区精华

风险与影响

关联 Issue

完整报告

参与讨论