基于CPU内存的v1 Transfer Connector

发表于 更新于 Disqus：

vLLM最近支持了外部加载Transfer Connector，基于LMCache给出的StorageSharedConnector的例子，我尝试实现了一个基于共享内存的Transfer Connector。

v1的接口是一个可以layer wise的实现。

实际上使用下来其实没有明确的区分Producer和Consumer的角色，在P2P的场景下可能比较明显，实际上谁生产kv cache谁消费kv cache其实没有明确规定。

这个的好处是Prefix Cache和KV Cache Transfer没有明确的区别了，Prefill和Worker之间唯一的区别就变成了max_token=1和max_token为真实值的区别了。

设想几个场景，Worker即是生成者，也可以是消费者，Prefill也可以即是生产者又是消费者：

  1. Worker 生成的对话可以存入一个中心化的缓存当中，在多轮对话的时候Prefill可以直接复用这个缓存，只需要计算新的用户对话。
  2. Prefill 基于新的对话可以生成一个缓存，被Worker使用，也可以被其他的Prefill在新的多轮对话中使用。

有一个比较hack的查看调用栈的方法就是在对应的接口函数抛出异常让程序崩溃，就能在stack trace上看到函数调用的路径了。

当然这个接口也可以实现P2P，毕竟他提供了对应的wait接口，无非是等中心化的缓存是否就绪还是Prefill的直接传输是否就绪区别了，这在提供的接口列表当中可以看到。

实现一个Connector需要关注几个接口。

Worker side

Layer Wise

vllm/vllm/attention/layer.py中可以看到。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
def unified_attention(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    layer_name: str,
) -> torch.Tensor:
    wait_for_kv_layer_from_connector(layer_name)

    forward_context: ForwardContext = get_forward_context()
    attn_metadata = forward_context.attn_metadata
    if isinstance(attn_metadata, dict):
        attn_metadata = attn_metadata[layer_name]
    self = forward_context.no_compile_layers[layer_name]
    kv_cache = self.kv_cache[forward_context.virtual_engine]
    output = self.impl.forward(self, query, key, value, kv_cache,
                               attn_metadata)

    maybe_save_kv_layer_to_connector(layer_name, kv_cache)
    return output

  1. 每一层会有wait_for_kv_layer_from_connector调用connector.wait_for_layer_load。

  2. 在计算结束以后会有maybe_save_kv_layer_to_connector调用connector.save_kv_layer。

他们分别对应了decode和prefill，其中wait是同步的，save是异步的。

Model Wise

在vllm/vllm/attention/layer.py中

1
2
3
4
5
6
7
8
9
10
11
12
self.maybe_setup_kv_connector(scheduler_output)

model_output = self.model(
    input_ids=input_ids,
    positions=positions,
    intermediate_tensors=intermediate_tensors,
    inputs_embeds=inputs_embeds,
)

self.maybe_wait_for_kv_save()
finished_sending, finished_recving = (
    self.get_finished_kv_transfers(scheduler_output))

  1. connector.start_load_kv 来自于 maybe_setup_kv_connector，在decoder的model forward之前调用，用于异步启动kv的load。

  2. connector.wait_for_save 来自于 maybe_wait_for_kv_save，在prefill的model forward之后调用，用于整体的save是同步的。

其他的一些接口包括：

  1. get_finished返回对于给定的request ids对应的已经完成的sending和recving的request ids。

  2. register_kv_caches用于connector提前注册kvcaches，在初始化kvcache的时候调用，这个应该是NIXL需要这样直接读取整个的kvcaches的显存地址做RDMA和注册。

  3. bind_connector_metadata，model forward之前bind metadata，这个数据结构是Metadata的数据结构是实现者自由定义的，在start_load_kv之前调用。

  4. clear_connector_metadata，model forward之后clear。

换成prefiller和decoder的视角来看

Prefiller Connector

每一层调用save_kv_layer，这个可以是异步的，在model foward之后会调用wait_for_save保证kvcache被传输完，不然其中的kvcache可能会被之后的forward所覆盖。
clear_connector_metadata可以帮助清理这次forward相关的metadata。

Decoder Connector

bind_connector_metadata帮助设置forward相关的metadata。
每一层调用start_load_kv，这个可以是异步的，在model forward之前调用，在每一层forward之前wait_for_layer_load，这个是同步的。

Scheduler Side

  1. get_num_new_matched_tokens，基于传入的num_computed_tokens获取可以从外部加载的kvcache，这个是给scheduler用的，表明decoder需要加载的tokens。computed_token代表已经计算过kvcache的token.
    调度器要额外分配一个external_computed_tokens的slots给外部加载用并且把这部分也算在computed_token，然后在根据budget_token - computed_token分配new_token。

  2. update_state_after_alloc 在scheduler分配slots以后更新connector内部状态，比如用于告知connector是否要加载kvcache。

  3. build_connector_metadata用于构建connector metadata的相关输出，不能修改输入中的schedulerOutput。

  4. request_finished 在request结束，blocks free之前被调用，可以帮助connector触发相关回调。

上面的接口比如register_kv_caches，bind_connector_metada和clear_connector_metadata不一定要实现，可以把他们理解为一些初始化路径，计算路径上的调用hook，我们希望在相关的hook上处理一些东西就实现这些接口。

CPUMemorySharedConnector

我实现了一个基于共享缓存的实现，Prefill只生成缓存，Worker只消费缓存。
主要是通过用layer和tokens hash做key创建SharedMemory。
完整的项目在这里。