标签： LLM

2 行代码校验大模型（如DeepSeek-R1）权重文件下载完整性

很多人在 DeepSeek-V3/R1 爆火之后，都希望体验本地运行“满血版”模型。但是满血版模型的权重参数文件有 600 多个 G，光权重文件就拆成了 163 个。

当你受不了 HuggingFace 官网的下载速度，用其它方法或者渠道获得了权重文件后，怎么确认这些权重文件是完整无损坏的呢？

这里介绍一个最简单的方法，仅需要 2 行代码。

环境

前提 1，你已经 clone 了不含权重文件的模型 git 仓库。以 DeepSeek-R1 为例，通过下面命令可以仅 clone 代码文件到 DeepSeek-R1 目录下：

GIT_LFS_SKIP_SMUDGE=1 git clone https://huggingface.co/deepseek-ai/DeepSeek-R1

前提 2，你已经用某种方法下载好了权重文件。请将这些权重文件放到已 clone 的 git 仓库目录内，以 DeepSeek-R1 为例，就是将 163 个 *.safetensors 文件移动到 DeepSeek-R1 目录下。

你也可以不移动权重文件，那么你就需要在执行第 2 行命令前将 checksum 文件移动到权重文件所在目录。

第 1 行代码

获得所有官方权重文件的 sha256 checksum，并保存成一个标准的 checksum 文件。这行代码需要在 git 仓库目录下执行

git lfs ls-files -l | awk '{print $1"  "$3}' > large_files.sha256

这行命令输出的文件内容形如：

c2388e6b127ce6664e35c5e2529c3ce4bfc99f4f7fb6fa48e92b29ed5e4922af  model-00001-of-000163.safetensors
5f450c75da7eb897b74a092eee65df8bb115fce81cccd2bbaeb220bd97197875  model-00002-of-000163.safetensors
...
913177d9e0dfb228769e0a13a386c34b919dcbb32a430ce230979f53bf7ae5bc  model-00163-of-000163.safetensors

第 2 行代码

根据官方权重文件的 checksum，检查本地文件的完整性。这个检查的执行速度会非常慢，因为它需要为每个文件计算 sha256sum，然后再与 checksum 文件做比对。

sha256sum -c large_files.sha256

这行命令的输出形如：

model-00001-of-000163.safetensors: OK
model-00002-of-000163.safetensors: FAILED
...
model-00163-of-000163.safetensors: OK

如果所有行的输出都是 OK，那么恭喜你，所有权重文件都没有损坏；如果有某行输出为 FAILED，就代表该文件没有通过完整性校验，你需要重新下载它。

此方法对所有标记为 LFS 的文件均有效，并不仅限于 *.safetensors 文件，比如量化模型 *gguf 权重文件，也可以同样用此方法校验。

LLM PD 分离背后的架构问题

PD 分离（Prefilling Decoding Disaggregation）推理是指将大模型推理的预填充阶段（P）和解码（D）阶段分离，以减少预填充与解码相互之间的影响，以便对两个阶段分别进行优化，提升 GPU 硬件的利用率，并减少推理延迟的一种推理技术方案。

在 DistServe、Mooncake 等论文中介绍分离式架构之后，DeepSeek V3 的报告让大家更进一步意识到 PD 分离可能是影响成本和性能的关键技术。

vLLM 对 PD 分离已经有了一个 1P1D 的实验版本。除此之外的开源框架大多还都不支持，不过很多已经在计划、实现中了。但纵览这些实现、文章或者计划，可以看到 PD 分离的架构选型上有很多问题需要思考，我尝试列举一下：

一、PD 是否直连传输？或是否需要 KV Cache Store/Pool？

PD 直连就是预填充节点直接将 KV Cache 发送给解码节点，它的好处是延迟低。但也意味着在整个 batch 的计算过程中锁定了P、D 节点的对应关系，一旦解码节点出现了问题，比如压力过大、服务出错、传输阻塞，在重试时无法仅调度 D 节点，需要重新进行整个预填充、解码过程。在 prompt 较长时，或者在 PD 节点数不对等的场景下，例如 2 个 P 对应到 1 个 D，重调度意味着抛弃较长或者多个 prefill batch，重调度的沉没成本较高。

使用 KV Cache Store/Pool 是在 P 和 D 之间增加了一个中间存储，预填充节点先将 KV Cache 写到中间存储，解码节点从中间存储读。这样做数据会多传输一次，增加了延迟，也增加了一些复杂度。但好处是容错性更好，还有就是预填充阶段本身也可以利用这个中间存储做 Prefix Caching。

中间存储也会对其它一些架构变动的复杂度产生影响，参见下面问题 四 和 五。

目前来看，Kimi Mooncacke、vLLM 的下一步设计、阿里 RTP-LLM 都使用或者计划使用基于 KV Cache Store/Pool 的方案，DeepSeek V3 报告中没有提到这部分。

在一些计算配比均衡、故障风险较小的场景下，比如同机多卡之间的 PD 分离，PD 直连的方案也有其简单易部署的优势。

二、P/D 是否按层发送/接收 KV Cache？

预填充最简单的实现是预填充节点完成第一个 token 的生成后，将所有的 KV Cache 传输给解码节点，这也是 vLLM 当前的实现。但这样实现有个问题，因为 KV Cache 的规模有可能非常大（尤其是原始 MHA），一个 batch 的 KV Cache 可能会是 GB 级别，都放在计算完成后传输，传输的延迟开销会比较大。

Kimi Mooncacke 和阿里 RTP-LLM 都采取了按层传输的方案，这是利用了 LLM 多层计算的自然特性。在完成一层的计算以后，就将这一层的 KV Cache 发送出去。这样 KV Cache 的发送就呈流式，既能降低延迟，也能使数据的发送更平滑。还存在一个更显著的优势，是 KV Cache 占用显存的时间更短，在显存紧张的情况下显存效率更高。

但按层发送对推理引擎的修改显然更大。我还没有看到开源的实现，猜测按层发送的引入对推理引擎的优化应该会有一定的影响，这里可能还需要一些精巧的设计才能减少影响。另外，按层发送对于 PD 非直连的场景下，中间存储的实现也会显著更复杂，QPS * num_hidden_layers，考虑到连续性可能还需要存储预分配和 session 保持。

因此对于 MLA 这种 KV Cache 偏小的注意力实现，比如 DeepSeek V3 的 KV Cache 是 576B/token/layer，是否要做按层发送，也许要看一下实际收益。

解码阶段和预填充阶段有所不同。解码需要多次迭代，在第一次迭代实现按层解码也没太大意义，而且涉及到计算的编排，应该需要拿到所有层的 KV Cache 才会开始计算。而且解码的计算时间比较长，如果解码的计算能够掩盖接收的延迟，不一定非要实现按层接收。

解码时按层接收，对调度也有一定挑战。从时序上来说，先发请求给预填充，完成后再发请求给解码会更自然。同时请求预填充和解码，需要处理一些同步问题，比如预填充压力大、解码等 KV Cache 超时等等。比如像阿里 RTP-LLM，它会观测预填充的排队情况，当一个请求进入预填充执行阶段时，解码端开始启动显存申请。

三、First Token 怎么处理

通常来说，预填充的同时会顺便把第一个 Token 计算出来，但计算到 hidden states 还是 token id 需要做一个选择。

计算到 hidden states 的好处是，预填充节点完全不需要加载和计算 lm_head 参数。比如 DeepSeek V3 的 lm_head 参数量是 0.9B，如果计算到 hidden states，这部分参数就完全不需要加载了。vLLM 目前就是采取的这个方式，预填充除了需要发送 KV Cache 之外，还需要发送一个 hidden states，解码时引擎也需要能支持加载 hidden states 延续计算。

计算到 token id 的好处是，发送的数据量小。以 DeepSeek V3 为例，hidden states 7K，token id 4B，完全可以跟着控制面消息传输。解码时引擎处理也更简单，因为 token id 到 token 的 detokenizer 一般是 CPU 查表，不涉及 tensor 的特殊处理。阿里 RTP-LLM 看起来采用的是这个方案。

四、Prefiller 和 Decoder 是否能相互转换？

当到达请求的 prompt 长度有差异性的时候，预填充和解码就会出现压力的不均衡问题。因为整体的吞吐取决于 P 和 D 的全局资源利用，当 P 过载但 D 闲置，或者 P 闲置但 D 过载的时候，成本和性能都不是最优的。

所以就需要考虑在 P 和 D 之间做负载均衡，要么从整个节点层面直接切换 P 和 D 的角色，要么 P 和 D 节点能够承担一些混杂的请求，比如通过 chunked prefill。

这时候 P 和 D 是否直连对实现复杂度就有一些影响了，如果有中间存储的存在，通过 PD 转换做负载均衡的实现难度会降低很多。

五、Decoder 能填充 KV Cache 吗？

如果业务应用场景中会将生成的 context 也作为下一轮的输入，还可能需要考虑 Decoder 填充 KV Cache，用于下一轮的 prefix caching 复用。这时候，KV Cache Store/Pool 的存在，对流畅交互有比较大的意义。

六、KV Cache Store/Pool 的设计抉择

有别于我们通常的 KV 存储，由于 GPU、RDMA（IB、RoCE）、NVLink 新硬件的存在，KV Cache Store/Pool 的设计抉择点会非常多。

在存储上，有 VRAM、DRAM、NVMe SSD，要选择 KV Cache Store 使用哪些介质。虽然对于 MHA 来说，因为 KV Cache 太大，基于 SSD 存储并不现实，但是对于 MQA、MLA 来说，NVMe SSD 并不是不可用。

在通信上，有 TCP、NVLink、RDMA、GPU Direct RDMA、NVMe over RDMA。为了更高的性能，KV Cache Store 在数据面上可能要考虑使用更快、更直接的传输方法。但 RDMA 对数据访问的抽象比 TCP 复杂很多，TCP 就是一端发一端收，但 RDMA 很多是单边操作。比如数据从 A 机 VRAM 发送到 B 机 DRAM，可能有以下方法：

  • A 从 VRAM 复制到 DRAM 再写 B 的 DRAM
  • A 从 VRAM 复制到 DRAM 再让 B 读 A 的 DRAM
  • A 直接从 VRAM 复制到 B 的 DRAM
  • B 直接读 A 的 VRAM

如果再加上 NVMe over RDMA，那要考虑的东西就更多了。P 发送到 Store，D 从 Store 接收，到底要通过哪些模式支持，是需要思考的。目前来看，预填充节点更适合单边写到 Store，这样能减少状态传输，更快地释放显存，但如果预填充节点也要读 prefix cache，那情况可能反过来；解码节点可能更适合单边读 Store。

在分布式架构上，无论是做集群式的 KV Cache Store，还是单机 side-car 式的 KV Cache Store，都需要存储一些 meta，并且在 P、D 之间传输一些控制信息。学术界有一些完全基于 RDMA 实现的分布式 KV 数据库，但目前看复杂度还是比较高，也没有开源的实现。目前业界实现还是倾向于使用传统的 RPC 方式来传输控制信息，并且通过分布式技术方案做 meta 节点的一致性、可靠性设计。

在接口 API 上，KV Cache Store 比传统的 KV Store 要复杂一些。比如要支持写的时候分 layer 写，读的时候能读到连续的内容；还可能要支持队列式的读，写完的 layer 可以很快被读走。如果要支持 prefix caching，还存在 KV Cache 的链式关系，写的时候不仅要分 layer，还要分 page，读的时候也是。TP/SP 等并行计算机制，对 API 可能还会有一些额外的要求。

在数据结构上，如果希望从 VRAM 直接写 Store，减少一次复制，引擎本身的 KV Cache 数据结构就需要与 Store 的数据结构进行一定程度的对齐；如果希望同时兼做 prefix caching，那 store 的数据排布就要考虑相同 prefix 的 page 更接近，甚至共享。比如用 prompt 的所有 page 的 hash 组成 string，按前缀 range 分桶，桶内对相同前缀做 merge/引用等等，这在存储优化上会是一个挑战。

整体来看，PD 分离的实现上有很多架构问题需要抉择，目前还没有一个理想的架构方案，或许未来也会是根据不同场景有很多参数化的灵活配置。

DeepSeek V3 模型各子模块参数量精算

网上很多文章一般只提到 DeepSeek V3 模型的总参数量，很少有人分析各子模块的参数量。我试着让各 AI 根据配置计算一下，没有一个靠谱的，只能自己算了。（也许本文的内容后续会变成 AI 回答本问题的 RAG 养料）

下面是根据 DeepSeek V3 开源仓库 https://huggingface.co/deepseek-ai/DeepSeek-V3，对 DeepSeek V3 各子模块参数量进行的精算，在计算复杂的 TP、DP、EP 拆分时可以用作基数参考。如有错误，烦请评论指出。

嵌入层 Embedding

"vocab_size": 129280, // Token 字典大小
"hidden_size": 7168,

DeepSeek V3 的嵌入层参数量是：

129280 * 7168 = 926,679,040 (~0.9B)

MLA

"hidden_size": 7168,
"num_key_value_heads": 128,
"v_head_dim": 128,
"kv_lora_rank": 512,

"num_attention_heads": 128,
"q_lora_rank": 1536,

"qk_nope_head_dim": 128,
"qk_rope_head_dim": 64,

"num_hidden_layers": 61,

单层 MLA 中 Q 的 LoRA 参数量是：

7168 * 1536 + 1536 + 1536 * 128 * (128 + 64) = 48,760,320

单层 MLA 中 KV 的 LoRA 参数量是：

7168 * (512 + 64) + 512 + 512 * 128 * (128 + 128) = 20,906,496

单层 MLA 中 WO 的参数量是

128 * 128 * 7168 = 117,440,512

pre 和 post attention layernorm 的参数量是：

7168 * 2 = 14336

所以 DeepSeek V3 的 MLA 部分共 61 层的总参数量是：

(48,760,320 + 20,906,496 + 117,440,512 + 14336) * 61 = 11,414,421,504 (~11B)

MoE

"num_hidden_layers": 61,
"hidden_size": 7168,
"moe_intermediate_size": 2048,  // 路由专家 MLP 的中间维度
"n_shared_experts": 1,          // 共享专家数量
"n_routed_experts": 256,        // 路由专家数量
"first_k_dense_replace": 3,     // 前几层使用dense替换MoE
"intermediate_size": 18432,     // 前3层 (9*moe_intermediate_size)

每个专家的参数量是：

7168 * 2048 * 3 = 44,040,192

路由 Gate 的参数量是：

256 * 7168 + 256 = 1,835,264

前 3 层 dense（固定激活 8 路由专家），前 3 层参数量是：

44,040,192 * 9 * 3 = 1,189,085,184

后 58 层稀疏（动态激活 8 路由专家），后 58 层参数量是：

(44,040,192 * 257 + 1,835,264) * 58 = 656,569,547,264

所以 DeepSeek V3 的 MoE 部分的总参数量是：

1,189,085,184 + 656,569,547,264 = 657,758,632,448 (~657B)

每次计算激活 1 个共享专家，8 个路由专家，所以 DeepSeek V3 MoE 部分的激活参数量是：

44,040,192 * 9 * 61 + 1,835,264 * 58 = 24,284,510,720 (~24B)

Layer 维度

前 3 层是 dense，没有 gate，基于上面的计算，DeepSeek V3 前 3 层每层的参数量是：

(48,760,320 + 20,906,496 + 117,440,512 + 14336) + (44,040,192 * 9) = 583,483,392

后 58 层是 MoE 稀疏激活专家，基于上面的计算，DeepSeek V3 后 58 层每层的参数量是：

(48,760,320 + 20,906,496 + 117,440,512 + 14336) + (44,040,192 * 257 + 1,835,264) = 11,507,286,272

输出层

DeepSeek V3 输出层的 RMSNorm 和 Linear 参数量是：

7168 和 129280 * 7168 = 926,686,208 (~0.9B)

总参数量

核对一下 DeepSeek V3 总参数量是否为 671B：

583,483,392 * 3 + 11,507,286,272 * 58 + 926,679,040 * 2 + 7168 = 671,026,419,200 (~671B)

核对一下 DeepSeek V3 激活参数量是否为 37B：

11,414,421,504 + 24,284,510,720 + 926,679,040 * 2 + 7168 = 37,552,297,472 (~37B)

这个与 README_WEIGHT.md 中提到的 36.7B 不同，我还没找到计算错误的地方。我理解也许是考虑到 embedding 层只是查表，并不是矩阵乘，所以实际激活参数是：

11,414,421,504 + 24,284,510,720 + 926,679,040 + 7168 + 7168 = 36,625,625,600 (~36.6B)

PS （20250208）

DeepSeek 更新了 README_WEIGHT.md ，激活参数量修正成了 36.6B，也去掉了包含 0.9B 输入 embedding 的注释。

MTP

DeepSeek V3 MTP 的 ebedding 和输出 head 与主模型共享，enorm 和 hnorm 的权重是：

7168 + 7168 = 14336

eh_proj 线性变换的权重规模是：

7168 * 14336 = 102,760,448

增加了一层 hidden layer，即第 61 层：

(48,760,320 + 20,906,496 + 117,440,512 + 14336) + (44,040,192 * 257 + 1,835,264) = 11,507,286,272

加起来 DeepSeek V3 MTP 的总参数量是：

11,507,286,272 + 102,760,448 + 14336 = 11,610,061,056 (~11.6B)

DeepSeek V3 MTP 的激活参数量是：

11,610,061,056 - 44,040,192 * (256 - 8) + 926,686,208 * 2 = 2,541,465,856

这个规模比 README_WEIGHT.md 中提到的 11.5B 独立参数，和 2.4B 激活参数都略大一点。

PS （20250208）

DeekSeek 更新了 README_WEIGHT.md ，MTP 的激活参数量由 2.4B 改成了 1.5B，可能跟上面的激活参数一样，都减去了 embedding 层。但在我的计算里，这个应该是 1.6B :)，还是略有不同。

11,610,061,056 - 44,040,192 * (256 - 8) + 926,686,208 = 1,614,779,648 (~1.6B)

DeepSeek V3：AI 大模型 infra 基建新高度

AI 工程化

2021 年初，贾扬清在阿里云开始推广 AI 工程化这个概念，我非常认同。我在 21 年中做技术规划的时候，提出“AI 到生产力的转化，需要更高的工程化能力”，并且将 AI 工程化的实施总结为几大方向：

  • 语义索引场景化
  • 算力调度混合化
  • 模型研发标准化
  • 优化技术普惠化
  • 模型超大规模化
  • 架构系统智能化

我的 AI 工程化团队在这些方向上也取得了许多成果。

The AI Model

但 2022 年底 LLM 大流行以后，情况发生了一些变化。原因主要是 LLM 让 AI models 变成了 The AI model，虽然这个 model 很大，也多多少少有一些变种，但从工程实践的角度来看，它并不“复杂”。

很多时候，工程架构解决的是复杂性问题。

比如，TensorFlow、PyTorch、PaddlePaddle 这些训练框架简化了搭建和训练神经网络的复杂度，在一段时间内，各种结构的网络层出不穷，大部分都是依托这些框架来实现的。

而对于 LLM 来说，模型结构相对固定，虽然也使用了框架的一些外围能力，但是模型结构核心部分已经逐渐变成全手写以达成最佳性能，典型的实现包括 FlashAttention、TRT-LLM 等。

而且 LLM 的接口调用是自然语言，因而也变得极其简单，所有的 LLM 模型几乎可以使用同一套 API。

当时看起来 LLM 不需要太多的架构基建工作。

Prefix Caching 和 Kimi

我的这个认知在思考 prefix-caching 作用的时候，有了一些改变。

在《应该把 Prefix Caching 当作一种效果优化技术》这篇博客中，我提到 Prefix Cache Aware Scheduling 是一件非常值得做的事情。而且从 Kimi 发表的论文来看，他们已经在实践了，但其它的技术报告提到这些工程架构工作的不多。

DeepSeek V3

前几天 DeepSeek AI 发布了 DeepSeek V3 版本，我一边在吐槽这 670B 的模型参数太大，下载太慢，一边在阅读它的技术报告。结果发现他们在模型的部署上，玩得更高端，给了我一些新的震撼。

首先，prefilling 和 decoding 分开部署。prefilling 4 机 32 卡，decoding 40 机 320 卡。这样一来，我之前《LLM 推理优化 Continuous Batching 及其实现》这篇博客中提到的 Continuous Batching 就不再需要了。两阶段分开后，prefill 的计算过程（长度）是确定的，其算力利用是充分的，不再需要中间停下来插入新的请求。其实 prefilling 能够分开部署，跟 DeepSeek 以前的研究也是分不开的，DeepSeek V2 引入的 MLA 对 KV Cache 做了大幅度的低秩压缩，可以显著降低 KV Cache 从 prefilling 节点传递到 decoding 节点的带宽和延迟。

其次，MoE 专家分开部署。因为 MoE 专家的激活是 Token 级别的，也就是说每个 Token 会决定走哪个专家进行计算，分开部署就可能会带来负载均衡问题：有些专家太忙，有些专家太闲。DeepSeek V3 为了解决这个问题，还做了复杂的负载均衡策略。例如：快速识别较忙的专家，部署冗余的专家副本以承担压力；重新调整专家在不同节点的布局，尽量利用跨 GPU 带宽而减少跨节点带宽（因为 IB 比 NVLink 要慢）；单卡冗余部署多专家，但通过全局的路由计算来优化专家的动态激活数量。

DeepSeek V3 还实现了计算和通信重叠。为了掩盖分布式计算过程中进行集合通信时的开销，将计算任务分为微批。一个微批进行集合通信时，进行下一个微批的计算。

此外，DeepSeek V3 在推理时还将 TP（Tensor）、DP（Data）、SP（Sequence）、EP（Expert）不同维度的并行化融合到了一起。单拿出来一种并行化方法也许现在的框架还是支持的，但这些方法组合在一起，我怀疑目前也没有什么推理加速框架能直接处理。

从技术报告中揭露的这些细节可以看出，为了发挥出模型的极致性能，DeepSeek 在 AI 大模型的分布式部署上花费了很大的心思。这也让 DeepSeek V3 成为目前公开资料可以看到的最复杂、最精巧的大模型 infra 设计。

这些 idea 以前也许不是没有人想到，但是 infra 的演进是有很高研发和试错成本的。当 DeepSeek 将这些路走通以后，也许未来的很多大模型公司，大模型框架，都会往沿着这个方向继续演进。

应该把 Prefix Caching 当作一种效果优化技术

我在 4 月份写的博客《LLM 推理优化 Prefix Caching 及其实现》带来了不少流量，这篇博客在 Google 搜索“Prefix Caching”时排在比较靠前的位置，也给我的微信公众号“边际效应”带来了超过 100 个关注者，由此可以看到大家对这项技术的广泛兴趣。

新认知

最近关于 Prefix Caching 我又产生了一次认知升级：在一个 Eureka 时刻，我忽然领悟到 Prefix Caching 实际上可以是一种效果优化手段——但让其充分发挥效能需要跟长上下文语言模型 Long-Context LLM 和基于 Prefix Cache 的调度技术相结合。

RAG v.s. Super Long Domain Specific Prefixes

为了减少模型的幻觉，提升时效性，我们往往会采取 RAG 技术来增强模型。召回的结果本身会增加 Prompt 的长度，会加大 Context 的计算量，会影响模型的推理时长，因而召回结果的补充也是谨慎克制的，不会让 Prompt 变得非常长。

但是 Prefix Caching 优化技术给我们开了个口子：如果把信息放在 Prompt 的共享 Prefix 中，加长 Prompt 的代价就由计算和时延，转化到了存储。而存储代价可以通过 Cache 复用来摊薄，这就是一笔很划算的经济账。

如果把 RAG 召回的结果当作 Prompt 里的变量，那么 Prefix 就是 Prompt 里的常量，当增加变量信息的规模受到限制时，增加常量信息的规模也可以提升生成的效果。

举个例子：如果你做 K12 领域的模型，把 12 个年级的语文知识点都放在 Prefix 里，然后再叠加 RAG，然后回答用户语文方面的提问，肯定能更好地保证生成的效果。

所以，LLM 技术后续的一个新发展分支，也许会是超长特定领域专用的前缀设计。

Inference Time Scaling

最近讨论很多的一个技术方向，是从基于预训练提升效果的 Scaling Law 转向基于 reasoning at inference time 的 Scaling Law，使用更多的推理计算来获得更好的效果。

但这些 Scaling Law 还是基于“computing”的 Scaling，我认为也许会存在基于“memory”的 Scaling Law，即更长的共享 Domain Specific Prefix 带来更好的效果。因为这里的“memory”是计算的缓存，memory 的共享本质上是计算的共享。

Long Context Large Language Model

在无法共享计算的场景下，长上下文的大语言模型应用往往由于其计算成本或显著延迟而无用武之地。但如果基于 memory 的 Scaling Law work 的话，那长上下文大语言模型将成为必经之路。

这也许是 Moonshot 和 Gemini 早早投入百万 Token 级别长上下文模型的背后逻辑。

Prefix Cache Aware Scheduling

在没有理想的大一统 Prefix 之前（大概率也不可能），共享 Prefix 只能是 Domain Specific 才最合适。那就意味着会有很多个共享 Prefix，而 Prefix Cache 占用的空间又是不可忽视的，所以不可能在每张卡/每台机器上都存储所有的共享 Prefix。

这就要求在用户请求调度时，根据 Prefix 分配到不同的卡/机器/集群上。Prefix Cache 也需要能够根据请求的热度在不同的服务器间弹性扩散或者收缩，甚至在显存、内存、SSD 之间调度。其实 Kimi 的论文“Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving”就是在解决这个问题，只是我之前没有意识到这个技术的广泛重要性。

LLM 推理优化 Prefix Caching 及其实现

Prefix Caching

上文中提到，Prompt 计算（Prefill 阶段）和生成阶段的计算特性很不相同。为避免重复计算，所有框架的 prefill 阶段主要作用就是给迭代的生成阶段准备 KV Cache。但这些 KV Cache 仅仅是为单次生成式请求服务的，那很自然的一种想法就是，KV Cache 能不能跨请求复用？

在某些场景下，多次请求的 Prompt 可能会共享同一个前缀（Prefix），比如拟人 Agent 的人物设定，文档阅读理解时的文档内容等。这些情况下，很多请求的前缀的 KV Cache 计算的结果是相同的，像一般互联网服务的请求缓存一样，可以被缓存起来，给下一个请求复用。

限制

但 KV Cache 跟其它服务缓存不一样的地方是，它太大了，以至于（目前）很难通过 Redis/Memcache 这种分布式缓存服务存取。比如对 13B LLM 模型来说，在 FP16 精度下单 token 的 KV Cache 大约是 1MB，假设要缓存的前缀有 500 个 token（大约800多个汉字），那就是 500MB。一般来说，我们不会每次请求去从分布式系统里读取/传输 500MB 的缓存，甚至都不会每次请求从内存往显存中拷贝 500MB 的缓存，所以大部分情况下，prefix cache 都会放在显存里。

这也就意味着，如果你想命中 prefix cache，必须把相同 prefix 的请求发到同一张 GPU卡上才行。

实现

由于不是普遍需求，加上前面说的限制，prefix caching 作为一个加速特性，不是很受关注，一般也不是默认开启的。各框架的实现和配置略有差异，这里简单做下记录，便于回顾。

刚开始 vLLM 的实现是给 generate 接口增加一个 prefix_pos 参数，通过 prefix_pos 输入参数为每个请求指定 prefix 长度，为 prefix 建一个带淘汰的哈希缓存。后来觉得这样做使用上不够便利，升级成了自动前缀缓存，即将 prompt 的 kv cache 分成 block，然后为 block 建设 LRU 缓存机制，这样就不必在接口上使用 prefix_pos 指定哪部分是 prefix 了。自动前缀缓存功能默认是不开启的，开启的配置项为 --enable-prefix-caching 。

TensorRT-LLM 与 vLLM 后来的实现类似，也是实现了 block kv cache，配置项是 enableBlockReuse，默认也是不开启的。代码未开源，无法看到实现。

Lmdeploy 的 Python 和 TurboMind C++ 版本的 prefix caching 功能都已经有了 PR，但现在（20240425）看还没有合入主干。有意思的是它没有使用 hash block 对应 token_id 子串的所有 token_id 前缀然后组成哈希表的方式，而是用 hash 当前 block 对应的 token_id 子串然后组成 trie 树的缓存管理结构。默认的参数名与 vLLM 相同，也叫做 --enable-prefix-caching。

HuggingFace TGI 现在看起来还没实现 Prefix Caching 功能。

Prompt Caching

除了 Prefix Caching 这种比较直观的工程优化，现在也有一些研究在看 Prompt 的其它缓存机制。比如设计一种机制让 prompt 模块化，不仅可以复用 prefix，还能复用中间的部分；或者通过 query 的相似性复用其它 query 的 prompt。

但目前看实现上都过于复杂，比如第一种要求模型使用不连续的 poition_id，这样就可以插入 token，但这种方式对 attention 的计算机制有一定的影响，难以说明它对效果的影响。

LLM 推理优化 Continuous Batching 及其实现

原理

Continuous Batching 是 LLM 推理优化的一项技术，作为这篇文章的知识背景不再赘述，目前流传最广的参考资料是这篇：《How continuous batching enables 23x throughput in LLM inference while reducing p50 latency》。它也有中文翻译，感兴趣可以搜一下，先看看。

图片来自：Anyscale 博客

虽然这篇资料介绍了它的主要原理，尤其是这张简单易懂的图，但是实现与原理是存在差异的，因为工程实现要解决很多现实问题。

(Re)scheduling 开销

如原文所说，Continuous batching 还有个别名，叫做：batching with iteration-level scheduling，这里的 iteration 就是指一次 decode 计算。也就是说在每次 decode 的迭代过程中，做 batch 的调度调整。

但调度本身不是无代价的，它可能涉及到接收和处理新的输入请求，重新组织输入数据的形状，甚至各种状态的重新初始化，这些都需要消耗 CPU 时间。这也就意味着在这段时间里，GPU 是闲着的，GPU 没有得到充分利用。

所以在实现时，程序并不会真的每个 iteration 都做 scheduling，目前看到有两种做法：

  • 合理间隔调度。比如每 16 次 decode 计算后，检查一下是否有新的输入，以及是否有空闲的槽位，然后对 batch 做一次调度调整。这能够显著降低调度的开销（TGI、lmdeploy、vLLM）。
  • 排队比例调度。比如当前 batch 中有 10 个请求的 decode 正在进行，而排队中有 12 个请求，超过了排队比例 1.2，那么就启动一次调度调整（TGI）。

KV Cache 重读

如果真的像图中那样，每个生成 Token 的 decode iteration 与一个 prompt token 的计算对齐，那 KV Cache 的利用就会很糟糕。因为它们需要在 Global Memory 与 Shared Memory 之间反复搬运，而且每次搬进来以后只被使用一次。

这本质上是 prefill 阶段（prompt 计算）与生成阶段的计算特性不同导致的，prefill 阶段并不适合 one-by-one 的 token 计算。

所以在实现时，程序并不会真的做 prefill 和生成的 token 对齐调度。目前看到的调度方法有三种：

  • 在重调度时，如果有新的请求进来，那么将新请求的 prefill 计算和要进行的 decode 计算做合并（Orca、vLLM-prev）。
  • 在重调度时，如果有新的请求进来，那么先对新的请求做 prefill 计算，然后再合并所有进行中的请求做 decode 计算（TGI、vLLM、lmdeploy）。
  • 先根据 decode 耗时估算出来每次 decode 同时能做 prefill 的 token 数量，在重调度时，如果有新的请求进来，对新请求的 prefill 按照上面的估算进行分块，然后将分块的 prefill 和其它请求的 decode 计算融合在一起，一定程度上减少了 KV Cache 重读的次数，又避免先做 prefill 计算带来的 Token 生成延时增加（Sarathi-Serve+vLLM）。

可调优能力

LLM 推理服务往往用于生产环境，而生产环境面临的情况是复杂多样的。

  • 对于做阅读理解的应用来说，Prompt 可能会非常长，但生成的内容可能会非常短，开发人员可能会更追求吞吐；
  • 对于聊天应用来说，Prompt 可能较短，生成的内容也不会太长，开发人员可能会更追求延迟；
  • 对于创作类应用来说，Prompt 可能很短，生成的内容会更长，开发人员可能会更追求首 Token 延迟。

对 Continuous Batching 实现来说，就要求它调度策略尽量清晰，并且参数可调。所以更灵活的实现，未来可能会更受欢迎。

Logits of API-Protected LLMs Leak Proprietary Information

看到一篇挺有意思的论文，大开脑洞，没想到还能这么玩，做一下粗读的笔记。

论文

标题：《Logits of API-Protected LLMs Leak Proprietary Information》，链接： https://arxiv.org/pdf/2403.09539.pdf 。

假设条件

典型 LLM 需要将最后一个 Transformer 块输出的嵌入向量转成要输出的 Token，这一步往往通过一个线性变换加 softmax 获取那个最大概率的 tokenid。

线性变换的权重是一个 vocabulary size * hidden size 的矩阵，比如 llama2-7B 的词表大小是 32000，hidden size 是 4096，那么线性变换权重矩阵的尺寸就是 32000x4096。这个矩阵再与 4096x1 的嵌入向量相乘，得到的就是 32000x1 的 logits 向量，其中每一个元素对应着一个词表中的 token 作为最终输出的概率。

上面这只是假设，也许 GPT 使用的是一个非线性变换，那论文内容可能就不成立了。

数学原理

这个线性变换将一个 4096 维的向量映射到了一个 32000 维的向量，从线性代数的角度来看，这是一个低维向高维的映射，所以它肯定不是一个满射（onto mapping）。也就是说，这个映射的像空间（image）只是 32000 维实数空间的一个子空间（subspace），而且这个像空间的秩（rank）最多是 4096。

这意味着可以找到不多于 4096 个线性无关的基向量（basis），使得这个像空间的每一个元素都能表示为这些基向量的线性组合。假设能采集到 4096 个线性无关的输出 logits，那这些 logits 就构成了像空间的一组基向量。

反过来想，如果你不知道 LLM 的 hidden size，那么你可以通过采集足够多的输出 logits，以保证有足够多的线性无关的向量。然后对矩阵进行奇异值分解（singular value decomposition），可以通过非 0 的奇异值个数推导出矩阵的秩。这个秩应该接近于模型的 hidden size。

逆向恢复 logits

遗憾的是，很多模型的 API 并没有输出完整的 logits 矩阵，但幸运的是，OpenAI 的 API 支持输出最多 top 5 个 token 的 logprobs，并且支持 logit_bias 干预 token 的输出。那就给了反复通过 API 调用来逆向恢复 logits 向量的可能。

但是具体方法我没看，粗读嘛，知道能做到就行了，有用到的时候再看吧。还有另一篇文章《Stealing Part of a Production Language Model》分析了在没有 logit_bias 甚至没有 logprobs 时该如何恢复 logits，我也没看，记录下链接 https://arxiv.org/pdf/2403.06634.pdf 。

无法输出的 Token

这篇论文还介绍了很多其它应用，太长没有看。比较有意思的一个引用是，在将嵌入向量映射到 logits 的过程中，如果一个 token 的嵌入向量在其它 token 的嵌入向量组成的凸包的内部，它就永远不可能被输出。扫了一眼引用的论文，证明没看懂，大致意思是 softmax 权重矩阵的低秩特性导致了可能输出 token 的排列在线性变换后不会出现在子空间里？实话说我感觉不像是很严谨的数学证明。。。

在 LLM 时代我们是否还需要倒排索引？

近些年，EBR（基于文本嵌入向量的召回）的强大语义召回能力让基于属性索引的传统倒排索引技术黯然失色，即使对专业搜索引擎来说，EBR 的应用也是越来越广泛 [1,2,3] 。尤其在 LLM（大语言模型）激起的 RAG（检索增强生成）技术框架下，大部分人似乎已经忘记了倒排索引，向量数据库成为 RAG 的标配。

但在 LLM 时代倒排索引真的没有用武之地了吗？我尝试列一下自己的思考。

  • Embedding 向量缺少 ground truth，但是倒排有。你无法通过直接观察或者测量，来明确一个向量指代的一段文本是否一定包含某些信息。但是如果这段文本在某个 term 的倒排拉链里，你可以从一定程度上明确这段文本包含了一些相关信息。
  • term 命中也是一种 attention。在训练模型时，我们总是希望 LLM 能关注到 context 中应该关注的信息，而忽略其它无关内容。那跟用户问题或者指令中某些 term 相关的内容，应该需要更多的关注。其实也可以类比人类肉眼的信息查找，人们也总是会扫到某些关键词，然后再仔细阅读它的上下文篇章。
  • 不基于倒排做召回，仍可以用倒排做粗筛。倒排作为一种可以查询 term 命中的高效结构，对 EBR 也许可以起到补充作用。例如对于某些 EBR 效果不够理想，误召回概率较高的场景下，对得分比较低的文档用命中信息作一次粗筛，能显著提升送给模型的 context 质量，也能减少对 LLM 计算资源的浪费。
  • term 命中的位置信息和权重不再重要。对于 LLM 来说，它会自行关注到 context 中需要关注的信息，不再需要位置信息或者权重来指示文本中哪些部分更重要。也就是说，倒排只需要解答 term 在文本中是否出现的问题，而不需要回答出现几次、在哪里出现的问题。
  • 也许倒排不再用 term，而是 token。term 依赖于切词，有一定的语义含义，term 集合空间一般有百万甚至千万的量级。但现在 LLM 大部分使用 BPE（Byte-Pair Encoding）分词器，token 集合空间只有几万到十几万的量级。使用 token 将显著减少倒排链的数量而优化其性能，但 token 存在没有归一化、分词边界不对齐的问题，是否可用还有待验证。

参考

[1] Guo, Ruiqi, et al. "Accelerating large-scale inference with anisotropic vector quantization." International Conference on Machine Learning. PMLR, 2020

[2] Huang, Jui-Ting, et al. "Embedding-based retrieval in facebook search." Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining. 2020.

[3] Liu, Yiding, et al. "Pre-trained language model for web-scale retrieval in baidu search." Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining. 2021.