2024-03-19 闲话

昨天在 b 站上找到了一个 talk：大语言模型低延迟推理技术的新进展 by Hao Zhang UCSD

https://www.bilibili.com/video/BV1wu4m1w7XC/?spm_id_from=333.337.search-card.all.click&vd_source=a8a652edec8ff61277e89ca67e8738cb

首先介绍了一个叫做 vLLM 的工作，主要贡献是 “reduce memory fragmentation with paging”，是一个 LLM memory management mechanism 的 optimization。

Preliminary

比如现在有一台专门提供 Chat 功能的 LLM 服务器，每个访问都会在 GPU 上预留足够的内存来存储 4096(或者 2048 1024 无所谓) 个 kv 向量来应对用户的请求。（也就是 Inference 阶段）

显然经常用不了这么多空间，那么就会浪费。和操作系统类似，一共有 reservation/internal fragmentation/external fragmention 三种浪费。分别表示，为了正在进行的 request 而预留的多余空间；预留空间和实际使用的空间中的 delta；若干个 request 之间的空余内存但是不够任何询问的（对于“顺序安排内存，使用完毕回收 reservation 中未使用的内存”的方法，会在 external frag 上消耗严重）。

因为上述的问题，在之前最出色的 orca “操作系统”上内存利用效率高达不到 40%……我们希望给一台服务器给更多的用户提供服务（更高的 throughput），所以需要提升单位算力使用效率。

method

解铃还须系铃人，memory usage 的问题可以使用了 os 方向的 paging 技巧解决，具体想法就是把 GPU 的内存分块，给正在进行的 requeset 赋予 virtual memory，通过 block table 来进行运行内存和实际内存之间的映射。

这样成功将 reservation 和 external fragmentation 消除了，同时将 internal fragmentation 降至 block size/request，使用效率高达 90%+。

同时它还可以通过 sharing 降低内存消耗，比如在 copilot 场景下，我们需要生成若干条候选，这时候 “prompt” 就需要 sharing，但是整行整行的内容我们完全没必要 copy 若干遍，只要在 block table 中标注它的实际内存地址即可，最后不满一行的需要复制。而这样也将 memory usage 大大降低。

另外讲的一个情形是在端侧使用 LLM 进行推理（比如参数都已经下载到本地了），这时候 setting 是我们有固定的算力（一般比进行 conventional 的 Inference 所需要的算力要多），同时我们单纯希望这个东西跑得快。

做法是随机一组 $y_1\dots y_m$ 作为可能得输出，根据 $y_1\dots y_{i-1}$ 通过 $\arg\max p(y_i|y_{1,i-1},x)$ 得到 $y_i'$ ，此时我们认为如果 $y_i = y_i'$ ，那么就得到了正确的 token。不难发现 iterate 一次可以得到至少 $y_1$ 正确， $i$ 次可以得到至少 $y_1\dots y_i$ 正确。当所有点都不动的时候根据一些定理可以知道结果和逐个 token 生成的结果一样。

这个方法的 optimization 在于如果你某次得到了多个 $y_t'$ 正确那么就赚了，我不知道论文里面的实践效果如何。这个方法能否做到 parallel decoding 我还没想清楚，同时需要注意的是这里单步 iteration 比输出单个 token cost 高。但是 under 这个 setting，确实利用了你端侧一般过剩的算力，本质上是 FlOPS 换 decoding step，是个不错的方法。

【总结与展望】