执行摘要

本 PR 修复了 ops.swap_blocks_batch 在 legacy default CUDA stream 下因 cuMemcpyBatchAsync 拒绝该 stream 而崩溃的 bug。通过在 C++ 内核中增加 stream 可用性检查，使该场景自动降级为逐份拷贝，并新增单元测试验证两种 stream 模式下的正确性。修复影响小，已合入主干。

功能与动机

PR 标题和 body 表明，在 legacy default stream（handle 0 / cudaStreamLegacy）上调用 swap_blocks_batch 时，cuMemcpyBatchAsync 会返回 CUDA_ERROR_INVALID_VALUE 导致进程崩溃。关联 Issue #46345 提供了复现步骤。该修复确保 KV offload 的块交换操作在任意 stream 上下文下都能稳定执行。

实现拆解

1. 修改 CUDA 内核：在 csrc/libtorch_stable/cache_kernels.cu 的 swap_blocks_batch 函数中，原有的条件 if (batch_fn != nullptr) 直接进入批量 fast path。新增 usable_stream 判断，仅当 stream 非空且非 cudaStreamLegacy 时才走 fast path，否则走早已存在的逐份拷贝 fallback 循环。
2. 新增测试覆盖：在 tests/v1/kv_offload/cpu/test_swap_blocks_batch.py 中，创建辅助函数 _addrs（收集 Tensor 指针）和 _run_batch（构造随机数据并调用 ops.swap_blocks_batch 后校验）。两个测试用例分别验证 legacy default stream 和 dedicated stream 下的正确性；测试通过逐元素比较确认数据拷贝正确。
3. 改动范围小：仅 2 个文件，核心逻辑变更仅 6 行，不涉及其他模块或外部接口。

csrc/libtorch_stable/cache_kernels.cu

核心修复文件，修改 swap_blocks_batch 函数增加 stream 可用性检查，决定了快速路径 vs 降级路径的分流逻辑。

// csrc/libtorch_stable/cache_kernels.cu ( 关键片段 )
// ... 前面的 batch_fn 获取逻辑
​
auto batch_fn = []() -> BatchFn {
    // 从动态库或缓存中加载 cuMemcpyBatchAsync 函数指针
    // 找不到时返回 nullptr
    return reinterpret_cast<BatchFn>(fn_ptr);
}();
​
// cuMemcpyBatchAsync 拒绝 legacy default stream (handle 0 / cudaStreamLegacy)，
// 并返回 CUDA_ERROR_INVALID_VALUE。这里将 legacy stream 也导向每份独立拷贝的
// fallback，确保在任何 stream 上下文中都能正确完成拷贝。
const bool usable_stream = stream != nullptr && stream != cudaStreamLegacy;
​
if (batch_fn != nullptr && usable_stream) {
    CUmemcpyAttributes attr = {};
    // ANY 属性允许 DMA 引擎乱序预取源数据（无并发写入冲突时安全）
    attr.type = CU_MEMCPY_ATTRIBUTE_ACCESS_MODE;
    attr.access_mode = CU_MEMCPY_ATTRIBUTE_ACCESS_MODE_ANY;
    // 执行批量异步拷贝
    // cuMemcpyBatchAsync(...);
} else {
    // 逐份拷贝 fallback：遍历所有 descriptor 依次调用 cuMemcpy
    // 该路径在 legacy default stream 或 batch_fn 不可用时执行
    for (int i = 0; i < n; ++i) {
        CUDA_CHECK(cuMemcpy(reinterpret_cast<void*>(dst_ptrs[i]),
                            reinterpret_cast<const void*>(src_ptrs[i]),
                            sizes[i]));
    }
}

评论区精华

Issue #46345 中，reviewer noooop 要求作者将提交回滚到 68c592d 并 hold，以等待 CI 的强制合并流程。作者 Etelis 执行回滚。最终 noooop 和 orozery 先后 approve 并合并。未出现设计层面的争论。

风险与影响

• 风险：变更仅限制 cuMemcpyBatchAsync 的使用，降级路径一直存在且正确，故风险极低。默认流下性能可能略低于批量路径，但这是正确的降级行为。未涉及安全或分布式隐患。
• 影响：用户：修复了 legacy default stream 下的崩溃，提高了 KV offload 的兼容性。团队：快速合并，无后续维护负担。

关联脉络

• 直接关联 Issue #46345，该 issue 报告了 bug 并复现了问题，是本 PR 的驱动力。
• 本 PR 是 KV offload 模块中 swap_blocks_batch 的兼容性修复，完善了 v1 引擎的流处理能力。近期历史 PR 中 #45845 也涉及 KV cache 管理，但属于不同功能线。