从零开始设计 SGLang 的 KV Cache

正文

请到「今天看啥」查看全文

• b) kv cache的实际业务有多种， MHA，GQA， MLA，DoubleSparse，如何做业务的隔离？

2.1 cache池

内存池的定义，如果每次都是要使用内存的时候，才去申请，效率会很低。容易导致碎片化，带来管理的困难。
但我们可以提前申请一大块内存，需要的时候从这个内存池去拿就好了，SGLang这里也一样。

2.2 二级cache池

kv cache有这些自定义类型，还在不断增长，MHA，MLA，DoubleSpars cache，管理起来比较麻烦。

使用二级内存池，一级记录high level信息，跟具体业务隔离，二级各种类继承，根据需要来调用。

2.3 一级内存池：req_to_token_pool

跟踪每个请求使用的token位置，具体的kv cache在二级内存池。

key：第i个req对应req_to_token第i行，第i个req的第j个token对应req_to_token第i行第j列

value：从二级缓存池的allocater获取的token_to_kv的内存块id(定位)

2.3.1 代码

初始化代码

sglang/srt/model_executor/model_runner.py

核心函数， init_memory_pool

初始化内存池

self.req_to_token_pool = ReqToTokenPool(
            size=max_num_reqs + 1,
            max_context_len=self.model_config.context_len + 4,
            device=self.device,
            use_records=False,
        )

输入的最大的req请求量和max content len

实现代码

sglang\srt\mem_cache\memory_pool.py

2.3.2 提供的功能

分配内存块

self.req_to_token = torch.zeros(
(size, max_context_len), dtype=torch.int32, device=device
)

最大的req量级，content len

内存块的write，alloc，free，内存池的常见操作，这里不做赘述，大家自行看代码就好

2.3.3 实际操作代码

prepare_for_extend()函数非常合适来看一级内存是如何操作。

此函数为extend的执行做准备，分配一二级内存，更新一级内存池。 二级内存池，因为此时kv cache的值还没有计算好，只是分配，要在forward的过程中才会写入。

代码路径

请到「今天看啥」查看全文