FlexKV 分析(三)——缓存的异步读写操作

KVTaskEngine 继承自 KVTaskManager，负责管理 KV 缓存的异步读写操作。

KVTaskEngine 中实现的缓存操作相关接口：

get_async：提交异步 GET 任务：从下级存储（CPU/SSD/Remote）读取 KV 缓存到 GPU。返回 (task_id, return_mask)，其中 return_mask 指示哪些 token 命中缓存。
put_async：提交异步 PUT 任务：将 GPU 上的 KV 缓存写入下级存储。返回 (task_id, return_mask)。
get_match：先进行缓存匹配（不依赖真实 GPU 块地址），返回匹配结果；内部创建带有假 slot mapping 的任务，等待上层设置真实地址后再启动。
put_match：类似 get_match，用于 PUT 操作的预匹配。

GlobalCacheEngine.get 过程分析

GlobalCacheEngine.get 根据请求的 token 序列，查询各级存储（CPU、SSD、Remote）中的前缀缓存，生成一个 TransferOpGraph，描述从下级存储到 GPU 的数据搬移路径，并返回相关的回调信息。

整体流程图

text 复制代码

GlobalCacheEngine.get(request_id, token_ids, token_mask, slot_mapping, layer_num, layer_granularity, dp_id)
│
├─ 1. 输入校验与对齐
│   ├─ _check_input()：确保类型/维度正确
│   ├─ 确定 aligned_length（按 block 对齐）
│   ├─ 根据 token_mask 计算 block 范围 (start_idx, end_idx)
│   └─ 将 slot_mapping 转换为 GPU block IDs
│
├─ 2. 构建 SequenceMeta（token_ids 按 block 切分，生成哈希）
│
├─ 3. 根据是否启用远程存储选择路径
│   ├─ 若不启用远程 → _get_impl_local()
│   └─ 若启用远程   → _get_impl_global()
│
├─ 4. 对返回的 transfer_graph 添加虚拟完成操作
│   └─ add_virtal_op_for_mutiple_finished_ops(transfer_graph, finished_ops_ids)
│       → 返回最终的 transfer_graph 和 task_end_op_id
│
├─ 5. 构建 return_mask（标记哪些 token 将从缓存中获取）
│   └─ 范围：[start_idx + skipped]*tokens_per_block 到 [start_idx+skipped+num_gpu_blocks]*tokens_per_block
│
├─ 6. 绑定 DP 组：transfer_graph.bind_to_dp_group(dp_id)
│
├─ 7. 锁定需要保护的节点（node_to_unlock）
│   └─ 对每个 device_type，调用 cache_engine.lock_node()
│
├─ 8. 构建总体回调 callback（_transfer_callback）
│   └─ 传输完成后解锁节点、回收临时缓冲区
│
├─ 9. 构建操作级回调字典 op_callback_dict
│   └─ 每个操作完成后将对应节点标记为 ready
│
└─ 10. 返回 (transfer_graph, return_mask, callback, op_callback_dict, task_end_op_id)

_get_impl_local（无远程存储）详细流程

_get_impl_local

text 复制代码

_get_impl_local()
│
├─ 1. 匹配 CPU 和 SSD 的前缀
│   └─ match_local(sequence_meta) → cpu_matched_result, ssd_matched_result
│       (每个 result 包含：physical_blocks, num_ready_matched_blocks, num_matched_blocks, last_node...)
│
├─ 2. 根据 mask 裁剪结果块
│   cpu_matched_blocks = result.physical_blocks[:num_ready_matched_blocks][block_mask_start:block_mask_end]
│   ssd_matched_blocks = 同上
│
├─ 3. 计算 fragment 长度
│   fragment12_num_blocks = max(len(cpu_matched_blocks), len(ssd_matched_blocks))
│   fragment1_num_blocks = len(cpu_matched_blocks)
│   fragment2_num_blocks = max(len(ssd_matched_blocks) - len(cpu_matched_blocks), 0)
│
├─ 4. 若无需传输（fragment12_num_blocks == 0）→ 返回空图
│
├─ 5. 确定 GPU 目标块
│   fragment12_gpu_blocks = gpu_block_ids[:fragment12_num_blocks]
│
├─ 6. 处理 fragment2（SSD 中有但 CPU 中没有的部分）
│   ├─ 若启用 GDS：
│   │   └─ 创建 TransferType.DISK2D 操作（SSD → GPU）
│   │       → 直接放入 finished_ops_ids
│   │       → op_node_to_ready 记录 SSD 节点信息
│   └─ 否则（普通 SSD I/O）：
│       ├─ 从 CPU cache 中分配临时块：cpu_cache_engine.take(fragment2_num_blocks)
│       ├─ 若分配不足 → 回滚并返回空图
│       ├─ 创建 TransferType.DISK2H 操作（SSD → CPU）
│       ├─ 若 CPU 侧前缀完全匹配且 ready，则插入新块并记录节点；否则标记为待释放
│
├─ 7. 创建 H2D 操作（CPU → GPU）
│   └─ 源：fragment12_cpu_blocks（可能包含 fragment1 + 新分配的 fragment2）
│   └─ 目标：fragment12_gpu_blocks
│   └─ 添加依赖：若存在 DISK2H，则 H2D 依赖它
│
├─ 8. 收集需要解锁的节点（CPU 和 SSD）
│   └─ node_to_unlock[DeviceType.CPU] = (cpu_node_to_unlock, size)
│   └─ node_to_unlock[DeviceType.SSD] = (ssd_node_to_unlock, size)
│
├─ 9. 收集需要释放的临时 CPU 块（buffer_to_free）
│   └─ 通常是未插入索引的 fragment2 块
│
└─ 10. 返回 (transfer_graph, finished_ops_ids, node_to_unlock, op_node_to_ready, buffer_to_free, num_gpu_blocks)

match_local 调用底层提供的radix tree做前缀匹配。

python 复制代码

    def match_local(self, sequence_meta: SequenceMeta) -> Tuple[MatchResult, MatchResult]:
        cpu_matched_result = MatchResult()
        ssd_matched_result = MatchResult()
        if self.cpu_cache_engine:
            cpu_matched_result = self.cpu_cache_engine.match(sequence_meta)
        if self.ssd_cache_engine:
            ssd_matched_result = self.ssd_cache_engine.match(sequence_meta)

        return cpu_matched_result, ssd_matched_result

_get_impl_global（包含远程存储）流程

_get_impl_global 在本地基础上增加远程存储（Remote）作为第三级缓存。传输路径为：Remote → CPU → GPU，同时可能将 CPU 中的数据回写至 SSD。

text 复制代码

_get_impl_global()
│
├─ 1. 匹配 CPU、SSD、Remote 三个存储的前缀
│   └─ match_all(sequence_meta) → cpu, ssd, remote results
│
├─ 2. 裁剪各自的 matched_blocks
│
├─ 3. 计算 fragment 长度
│   fragment1_num_blocks = len(cpu_matched_blocks)
│   fragment12_num_blocks = max(len(cpu), len(ssd))
│   fragment123_num_blocks = max(len(cpu), len(ssd), len(remote))
│   fragment2_num_blocks = fragment12_num_blocks - fragment1_num_blocks
│   fragment3_num_blocks = fragment123_num_blocks - fragment12_num_blocks
│
├─ 4. 若无需传输 → 返回空图
│
├─ 5. 确定 GPU 目标块
│   fragment123_gpu_blocks = gpu_block_ids[:fragment123_num_blocks]
│
├─ 6. 处理 fragment2+3（SSD 和 Remote 中有但 CPU 中没有的部分）
│   ├─ 从 CPU cache 分配临时块：fragment23_cpu_blocks
│   ├─ 若分配不足 → 回滚并返回空图
│   ├─ 创建 DISK2H 操作（SSD → CPU，对应 fragment2）
│   └─ 创建 REMOTE2H 操作（Remote → CPU，对应 fragment3）
│
├─ 7. （可选）将 Remote 读取的数据同时写入 SSD（写回策略）
│   └─ 条件：enable_ssd && remote2h 存在 && SSD 前缀完全匹配
│       ├─ 从 SSD cache 分配块
│       ├─ 创建 H2DISK 操作（CPU → SSD），并依赖 REMOTE2H
│       └─ 插入 SSD 索引
│
├─ 8. 创建 H2D 操作（CPU → GPU），依赖 DISK2H 和 REMOTE2H
│
├─ 9. 收集需要解锁的节点（CPU, SSD, Remote）
│
├─ 10. 收集需要释放的临时 CPU 块（若有未插入索引的块）
│
└─ 11. 返回结果

远程数据先拉到 CPU 固定内存，再转到 GPU；同时可写回 SSD 以加速未来访问。

使用 protected_node 参数在分配块时防止正在匹配的节点被驱逐。

get_async 处理过程

KVTaskEngine.get_async

text 复制代码

get_async(token_ids, slot_mapping, token_mask, layer_granularity, dp_id, task_id)
│
├─ 1. 参数预处理
│   ├─ 若 token_mask 为 None，设为全 True（长度与 token_ids 相同）
│   ├─ 若 layer_granularity == -1，设为 model_config.num_layers
│   └─ 若 task_id == -1，调用 _gen_task_id() 生成新 ID
│
├─ 2. 创建任务（_get_match_impl 内部）
│   │
│   ├─ 调用 create_get_task(...)
│   │   ├─ 调用 cache_engine.get(...)
│   │   │   ├─ 输入：task_id, token_ids, token_mask, slot_mapping, num_layers, layer_granularity, dp_id
│   │   │   └─ 输出：TransferOpGraph, return_mask, callback, op_callback_dict, task_end_op_id
│   │   ├─ 创建 KVTask 对象
│   │   │   ├─ task_type = GET
│   │   │   ├─ status = READY (因为 is_fake_slot_mapping=False)
│   │   │   ├─ graph = 返回的 TransferOpGraph
│   │   │   ├─ return_mask = 返回的 mask
│   │   │   ├─ callback, op_callback_dict 等
│   │   │   └─ task_end_op_id
│   │   ├─ 存入 self.tasks[task_id]
│   │   └─ 记录映射：graph.graph_id -> task_id
│   │
│   ├─ 处理空图：_process_empty_graph(task_id)
│   │   └─ 若 graph.num_ops == 0，直接调用 _mark_completed(task_id)
│   │       └─ 状态变为 COMPLETED，触发回调
│   │
│   └─ 返回 (task_id, return_mask)
│
├─ 3. 追踪记录（tracer.trace_request）
│   └─ 记录请求类型="GET"、task_id、token_ids、slot_mapping 等
│
├─ 4. 启动任务：_launch_task(task_id)
│   │
│   ├─ 获取任务对象
│   ├─ 检查状态：若已完成或状态不是 READY → 返回或报错
│   ├─ 设置状态 = RUNNING
│   ├─ 获取 graph
│   ├─ 若 graph.num_ops > 0:
│   │   └─ 遍历所有 transfer_handles（process/thread/remote 模式）
│   │       └─ 对每个 handle 调用 handle.submit(transfer_graph)
│   │           └─ 提交到底层 TransferEngine（异步执行）
│   └─ （若 graph.num_ops == 0，已在第2步标记完成）
│
└─ 5. 返回 (task_id, return_mask) 给调用者

_launch_task 主流程（从任务启动到传输引擎提交）

KVTaskManager._launch_task

text 复制代码

_launch_task(task_id)
│
├─ 获取任务对象 task = self.tasks[task_id]
├─ 检查状态：若已完成或非 READY 则返回/报错
├─ 设置 task.status = RUNNING
├─ 获取 transfer_graph = task.graph
│
└─ 若 graph.num_ops > 0:
    │
    └─ 遍历 self.transfer_handles (List[TransferManagerHandle])
        │
        └─ 对每个 handle 调用 handle.submit(transfer_graph)
            │
            ├─ [thread 模式] TransferManagerIntraProcessHandle.submit
            │   └─ self.transfer_manager.submit(transfer_graph)
            │       └─ TransferManager.submit → transfer_engine.submit_transfer_graph(graph)
            │
            ├─ [process 模式] TransferManagerInterProcessHandle.submit
            │   └─ 通过 Pipe 发送 {'type':'submit', 'transfer_graph': graph}
            │       └─ 子进程 _process_worker 接收 → transfer_manager.submit → 同上
            │
            └─ [remote 模式] TranserManagerMultiNodeHandle.submit
                └─ 通过 ZMQ PUSH 发送 {'type':'submit', 'graph': graph, ...}
                    └─ 远程 TransferManagerOnRemote._polling_worker 接收
                        └─ 调用 self.submit(graph) → 本地 TransferManager.submit → 同上

最终统一进入 TransferEngine.submit_transfer_graph(transfer_graph)

三种 TransferManagerHandle 模式对比:

text 复制代码

┌─────────────────────────────────────────────────────────────────┐
│                    TransferManagerHandle.submit                 │
└─────────────────────────────────────────────────────────────────┘
                │
                ├── mode="thread" ──► TransferManagerIntraProcessHandle
                │                     │
                │                     └─ 直接调用本进程 TransferManager.submit
                │                         └─ TransferEngine.submit_transfer_graph
                │
                ├── mode="process" ──► TransferManagerInterProcessHandle
                │                     │
                │                     ├─ 通过 Pipe 发送命令到子进程
                │                     └─ 子进程接收后调用 TransferManager.submit
                │
                └── mode="remote" ──► TranserManagerMultiNodeHandle
                                      │
                                      ├─ 通过 ZMQ PUSH 发送到远程节点
                                      └─ 远程 TransferManagerOnRemote 接收后处理

TransferManagerIntraProcessHandle

TransferManagerInterProcessHandle

TranserManagerMultiNodeHandle

TransferEngine 内部提交与调度流程

TransferEngine.submit_transfer_graph

TransferEngine._scheduler_loop

text 复制代码

TransferEngine.submit_transfer_graph(graph)
│
└─ self.task_queue.put(graph)      # 非阻塞，立即返回

[后台调度线程 _scheduler_loop 独立运行]
│
while _running:
│
├─ 1. 从 task_queue 取出所有新图
│   └─ for graph in task_queue:
│       └─ scheduler.add_transfer_graph(graph)   # 解析图的 DAG
│
├─ 2. 从 finished_ops_queue 收集已完成的操作
│   └─ while 队列非空:
│       ├─ op_id = finished_ops_queue.get()
│       ├─ op = op_id_to_op[op_id]
│       ├─ free_op_from_buffer(op, pin_buffer)   # 释放 pin memory 槽位
│       ├─ completed_queue.put((op.graph_id, op.op_id))
│       ├─ finished_ops.append(op)
│       └─ 删除 op_id 映射
│
├─ 3. 调度器生成下一步可执行操作
│   └─ completed_graph_ids, next_ops = scheduler.schedule(finished_ops)
│
├─ 4. 将 next_ops 分发给对应 Worker
│   └─ for op in next_ops:
│       ├─ 若 transfer_type == VIRTUAL:
│       │   └─ completed_queue.put((op.graph_id, op.op_id))
│       ├─ 否则:
│       │   ├─ op_id_to_op[op.op_id] = op
│       │   ├─ register_op_to_buffer(op, pin_buffer)   # 分配 pin memory 槽位
│       │   └─ _assign_op_to_worker(op)
│       │       └─ 根据 op.transfer_type 选择 Worker 进程/进程组
│       │           └─ worker.submit_transfer(op)   # 放入 Worker 的输入队列
│
└─ 5. 处理完成整个图的通知
    └─ for graph_id in completed_graph_ids:
        └─ completed_queue.put((graph_id, -1))

（循环间隔 0.001 秒，避免空转）

TransferEngine._init_workers 创建许多worker：

gpucpu_workers，cpussd_read_worker，remotecpu_read_worker，remotecpu_write_worker，gds_workers。

worker执行从TransferEngine中获取的op。完成的op，放入finished_ops_queue。

_assign_op_to_worker 根据op.transfer_type的类型，选择对用的worker。

op.transfer_type是GlobalCacheEngine根据缓存的命中情况确定的。

Worker 执行与完成通知流程

以GPUCPUTransferWorker为例。

text 复制代码

Worker 进程（独立 multiprocessing.Process）
│
├─ 初始化：
│   ├─ 注册 CUDA 设备、创建流
│   ├─ 建立与主进程的队列连接（input_queue, finished_ops_queue）
│   └─ 设置 ready_event，通知主进程已就绪
│
├─ 主循环：
│   while True:
│       ├─ op = input_queue.get()    # 阻塞等待操作
│       ├─ 若 op 为 None 则 break (关闭信号)
│       │
│       ├─ 根据 transfer_type 执行传输：
│       │   ├─ H2D: CPU block → pin buffer → GPU block
│       │   ├─ D2H: GPU block → pin buffer → CPU block
│       │   └─ （同步：cudaStreamSynchronize 或使用事件）
│       │
│       └─ finished_ops_queue.put(op.op_id)   # 通知主引擎
│
└─ 关闭：释放 CUDA 资源，退出进程

数据传输由GPUCPUTransferWorker._transfer_impl 调用 transfer_kv_blocks 触发。

transfer_kv_blocks

transfer_kv_blocks的pybind接口

cpp 复制代码

m.def("transfer_kv_blocks", &transfer_kv_blocks_binding,
      "Transfer multi-layer KV-cache between CPU and GPU",
      py::arg("gpu_block_id_tensor"), py::arg("gpu_tensor_ptrs_tensor"),
      py::arg("gpu_kv_stride_in_bytes"), py::arg("gpu_block_stride_in_bytes"),
      py::arg("gpu_layer_stride_in_bytes"), py::arg("cpu_block_id_tensor"),
      py::arg("cpu_tensor"), py::arg("cpu_kv_stride_in_bytes"),
      py::arg("cpu_layer_stride_in_bytes"),
      py::arg("cpu_block_stride_in_bytes"), py::arg("chunk_size_in_bytes"),
      py::arg("start_layer_id"), py::arg("num_layers"),
      py::arg("transfer_num_cta") = 4, py::arg("is_host_to_device") = true,
      py::arg("use_ce_transfer") = false, py::arg("is_mla") = false,
      py::arg("gpu_block_type") = 0);

transfer_kv_blocks_binding调用 transfer_kv_blocks。

提供了两种数据拷贝方法：

Copy Engine using cudaMemcpyAsync
Custom kernel transfer

上层完成事件轮询与任务状态更新

KVTaskManager._update_tasks

text 复制代码

KVTaskManager._update_tasks(timeout)
│
├─ 调用 transfer_handle.wait(timeout) 获取完成事件
│   └─ 内部最终从 TransferEngine.completed_queue 读取
│       └─ 返回 List[Tuple[graph_id, op_id]]   # 其中 op_id=-1 表示整个图完成
│
└─ 遍历每个完成事件 (graph_id, op_id):
    │
    ├─ 通过 graph_id → task_id 映射找到任务
    │
    ├─ 若 op_id == -1:
    │   └─ _mark_completed(task_id)   # 设置状态为 COMPLETED，调用 callback
    │
    ├─ 若 op_id == task.task_end_op_id:
    │   └─ 标记 task.task_end_op_finished = True
    │
    └─ 若 op_id 在 task.op_callback_dict 中:
        └─ 调用对应的 op 回调函数

（上层接口 wait() / try_wait() 会循环调用 _update_tasks 并检查完成条件）

_get_completed_ops

text 复制代码

_get_completed_ops(timeout)
│
├─ results = []                     # 最终返回的完成事件列表
│
├─ 对于每个 transfer_handle in self.transfer_handles:
│   │
│   ├─ completed_ops = handle.wait(timeout)
│   │   └─ 返回 List[(op_id, graph_id)]，其中 op_id = -1 表示整个 graph 完成
│   │
│   └─ 对于每个 (op_id, graph_id) in completed_ops:
│       │
│       ├─ 如果 op_id == -1：     # 整个 graph 完成
│       │   ├─ count = self.uncompleted_graphs.get(graph_id, 0) + 1
│       │   ├─ 如果 count == self.required_completed_count:
│       │   │   ├─ results.append((-1, graph_id))
│       │   │   └─ 删除 self.uncompleted_graphs[graph_id]
│       │   └─ 否则: self.uncompleted_graphs[graph_id] = count
│       │
│       └─ 否则：                 # 单个操作完成
│           ├─ count = self.uncompleted_ops.get(op_id, 0) + 1
│           ├─ 如果 count == self.required_completed_count:
│           │   ├─ results.append((op_id, graph_id))
│           │   └─ 删除 self.uncompleted_ops[op_id]
│           └─ 否则: self.uncompleted_ops[op_id] = count
│
└─ 返回 results

TransferManagerIntraProcessHandle.wait 处理流程分析

text 复制代码

TransferManagerIntraProcessHandle.wait(timeout)
   └─ self.transfer_manager.wait(timeout)   # transfer_manager = TransferManager 实例
        └─ self.transfer_engine.get_completed_graphs_and_ops(timeout)

TransferEngine.get_completed_graphs_and_ops

text 复制代码

TransferEngine.get_completed_graphs_and_ops(timeout)
│
├─ 1. 检查本地完成队列 self.completed_queue 是否为空
│   └─ 若 empty() → 直接返回 []（非阻塞快速路径）
│
├─ 2. 尝试获取第一个完成事件（带可选超时）
│   └─ first = self.completed_queue.get(timeout=timeout)
│       ├─ 如果队列在 timeout 时间内仍为空 → 抛出 queue.Empty
│       │   └─ 被捕获后返回 []
│       └─ 否则得到 (graph_id, op_id) 元组
│
├─ 3. 将 first 加入结果列表
│   └─ results.append(first)
│
├─ 4. 非阻塞地取出队列中剩余的所有完成事件
│   └─ while not self.completed_queue.empty():
│       └─ results.append(self.completed_queue.get_nowait())
│
└─ 5. 返回 results (List[Tuple[int, int]])

与调度线程的协作（填充方）

TransferEngine._scheduler_loop

text 复制代码

[后台调度线程 _scheduler_loop]
│
├─ 从 finished_ops_queue 获取已完成的操作 op
├─ 对每个 op：self.completed_queue.put((op.graph_id, op.op_id))
│
├─ 当调度器返回 completed_graph_ids（整个图完成）：
│   └─ for graph_id in completed_graph_ids:
│       └─ self.completed_queue.put((graph_id, -1))

completed_graph_ids, next_ops = self.scheduler.schedule(finished_ops)

只有当TransferOpGraph._op_map中所有op完成时，schedule才会返回completed_graph_ids。_scheduler_loop才会执行：completed_queue.put((graph_id, -1))

python 复制代码

class TransferOpGraph:
    _next_graph_id = 0
    _lock = threading.Lock()

    def __init__(self) -> None:
        self.graph_id = self._get_graph_id()
        self._op_map: Dict[int, TransferOp] = {}
        self._ready_ops: Set[int] = set()
        self._trigger_ops: Set[int] = set()
        self._gpu_transfer_op_id: List[int] = []