执行摘要

• 一句话：修复 packed KV cache 卸载时 set_ 溢出
• 推荐动作：该 PR 值得快速合并。改动简单、目的明确、风险低，且解决了生产环境中实际的崩溃问题。可以注意未来是否需要在 as_strided 前后增加边界断言以提高安全性。

功能与动机

当注意力层被打包到共享 KV cache 张量中（例如 MLA 模型使用非均匀页面大小），每个层的 page_size 可能小于 block_stride。torch.Tensor.set_(storage, offset, size, stride) 会验证 offset + size_bytes <= storage_size，但对于 packed 布局，最后一个 block 的 offset + page 会延伸到下一层区域，虽然是有效内存，但会触发 set_ 的边界检查失败。

实现拆解

1. 创建原始存储张量：先创建一个空的 int8 张量，通过 set_(layer_kv_cache.untyped_storage()) 将其绑定到共享的 untyped storage 上，获得一个覆盖整个存储的原始视图。
2. 替换 set_ 为 as_strided：使用 torch.as_strided(raw, (num_blocks, page), (block_stride_bytes, 1), byte_offset) 在原始张量上创建子视图，该视图具有与原来 set_ 相同的 shape、strides 和 storage_offset，但绕过了 set_ 的严格边界检查。
3. 保持 MambaSpec 分支不变：对于 MambaSpec 分支，由于 first_state_tensor 的 storage_offset 为 0，原有逻辑仍然正确，无需修改。

关键文件：

• vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py（模块 KV 连接器；类别 source；类型 core-logic；符号 register_kv_caches）: 单文件变更，核心逻辑：将 set_ 替换为 as_strided 修复 packed KV cache 视图创建时的边界检查溢出。

关键符号：register_kv_caches

关键源码片段

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

单文件变更，核心逻辑：将 set_ 替换为 as_strided 修复 packed KV cache 视图创建时的边界检查溢出。

# vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
# 在 register_kv_caches 方法中，针对 AttentionSpec 分支的修改
​
page = layer_kv_cache_spec.page_size_bytes
elem_size = layer_kv_cache.element_size()
byte_offset = layer_kv_cache.storage_offset() * elem_size
block_stride_bytes = (
    layer_kv_cache.stride(0) * elem_size
    if layer_is_packed[layer_name]
    else page
)
# 创建一个覆盖整个 untyped storage 的原始 int8 张量
# 使用 set_ 将空张量绑定到共享存储，但不指定 offset/size/stride
raw = torch.empty(
    0,
    dtype=torch.int8,
    device=layer_kv_cache.device,
).set_(layer_kv_cache.untyped_storage())
# 使用 as_strided 在原始张量上创建子视图，绕过 set_ 的严格边界检查
# 对于 packed 非均匀 KV cache，每个层的 page_size 可能小于 block_stride
# set_ 会验证 offset + size_bytes <= storage_size，但 packed 布局中
# 最后一个 block 的 offset + page 会延伸到下一层区域（合法内存）
tensors_per_block[layer_name] = (
    torch.as_strided(
        raw,
        (num_blocks, page),
        (block_stride_bytes, 1),
        byte_offset,
    ),
)

评论区精华

无 review 讨论。

• 暂无高价值评论线程

风险与影响

• 风险：风险较低。改动仅限 register_kv_caches 方法中 AttentionSpec 分支的 8 行代码，且已用 torch.as_strided 替代 set_，语义等效但绕过边界检查。潜在风险是：如果底层的 untyped storage 确实小于视图范围（例如配置错误导致 byte_offset + (num_blocks-1)*block_stride_bytes + page 超出 storage），则 as_strided 不会报错，后续访问可能读到越界内存，导致难以调试的损坏。但这种情况在正确配置下不会出现，且原有 set_ 也会对合法的 packed 布局拒绝。MambaSpec 分支未修改，无影响。
• 影响：影响范围：仅限于启用 KV offloading、且使用 packed non-uniform KV cache（如 MLA 模型）的用户。修复后这些场景不再因 set_ 溢出而崩溃。改动量小，无 API 变更，无性能影响。已在 GLM-5.2-FP8 生产环境验证。
• 风险标记：边界检查绕过, 缺少测试覆盖

关联脉络

• 暂无明显关联 PR