从transformer到PD分离

PD分离是一种大模型推理加速技术，且只能用在decoder-only架构的模型上。它将推理时的prefill部分和decoder部分分别在不同硬件上完成，加快模型推理的速度。

transformer结构

目前，AI技术已经毫无争议的成为了人类目前最重要的发明之一。大语言模型（LLM），多模态模型，扩散模型（diffusion）等大模型成为人们创造新应用的神兵利器，而目前所有大模型的底层结构全都要归于transformer。

Encoder-Decoder

Encoder-Decoder结构是transformer模型最原始的结构，随着transformer一同被提出，最开始用作文字翻译。其结构可由下图表示。其中左侧红框为encoder，右侧蓝框为decoder。推理的大致流程为数据从左侧输入到encoder中，之后encoder输出中间矩阵，作为右侧decoder的输入之一。decoder做自回归，即将该时间步的输出作为下一事件时间步的输入。推理过程如下图。
最开始，原始的句子通过encoder生成中间表达，输入到右侧decoder，decoder输出第一个翻译出的单词。实际上decoder有两个输入，在推理刚启动时，我们会向decoder输入一个特殊标记，向模型说明这是一个句子的开始。
之后decoder开始自回归，将上次输出的结果加上encoder的中间表达作为这次的输入，继续自回归。

可以将encoder作为decoder翻译的参考，就像我们人类翻译一样，翻译必须需要原始文本才能做翻译，encoder的作用就是将人类的自然语言转化成模型可以理解的语言。

Decoder-Only

但是现实生活中很多NLP的任务是没有原始文本的，比如简单的续写句子，模型没有原始文本可以参考，encoder就没有作用了，因此发展出了decoder-only模型。

decoder-only模型的推理过程和encoder-decoder模型中decoder的推理过程别无二致，同样是模型不断地做自回归。但是decoder推理时需要两个输入，没有了encoder，中间矩阵如何给出？答案是不要中间矩阵，直接通过给出的句子继续自回归。我们以续写句子这个任务为例，假设我给出“我是”，要求模型从“我是”后面续写完整的句子，模型可能会给出“我是中国人”这样的答案，我们知道模型需要一个token一个token的输出，我们假设一个字就是一个token，模型的输出情况会是下面这样：

我是中
我是中国
我是中国人

我们可以看到自“我是”之后，模型每次自回归就会输出一个token，那么我们就可以认为“我是”是模型已经输出过的token，我们就可以将“我是”提前编码好，输入到decoder中，推动decoder开始自回归，这也就是为什么提示词会被称为prompt。

Mask

我们回看上面的推理过程，发现每一次自回归都使用的之前的输入，那么我们能否将之前的输出保存起来，这样进行下一次自回归的时候就不需要重新计算之前输出的token了。当然可以，但在这之前，我们需要深入模型，了解transformer是如何处理数据的。

数据变成token，经过词嵌入处理后，会进入自注意力模块（self-attention）。每个词嵌入向量会与三个随机权重的矩阵相乘，分成查询向量，键向量，值向量三个向量，多个词向量连在一起与随机权重矩阵相乘，就形成了查询矩阵Q，键矩阵K，值矩阵V。之后将Q和K相乘，除以一个常量（这个常量是为了不要让数太大，稳住梯度），计算softmax之后乘上V就算出了attention，最后，通过某些算法得出最有可能的词。

在实际推理过程中，计算Q和K的结果后，还会引入mask（事实上在训练过程中也有这个mask，但与推理时的目的不同），具体而言，就是将矩阵与一个下三角矩阵按位相乘，使矩阵也变成一个下三角矩阵。

KV Cache

将整个推理的过程数学化。回想整个推理过程，当第一个token到来时，token会被分成Q，K，V三个向量（假设没有prompt），通过下面的式子计算出结果。

在第二个token到来时，与第一个token的结果组成一个矩阵，分成Q，K，V后，QK的结果经过mask，通过softmax之后与V相乘得到结果。这里mask在softmax之前和之后进行都可以，只要与V相乘的矩阵是一个下三角矩阵就行。

其中就是模型推理出的第二个token的结果。

根据上面的式子，我们可以继续写出和的表达式，
$$
\begin{align}
\operatorname{Att}{1}(Q, K, V) & = \operatorname{softmax}\left(Q{1} K_{1}^{T}\right) \overrightarrow{V_{1}} \
{A t t_{2}}(Q, K, V) & = \operatorname{softmax}\left(Q_{2} K_{1}^{T}\right) \overrightarrow{V_{1}}+\operatorname{softmax}\left(Q_{2} K_{2}^{T}\right) \overrightarrow{V_{2}} \
\operatorname{Att}{3}(Q, K, V) & = \operatorname{softmax}\left(Q{3} K_{1}^{T}\right) \overrightarrow{V_{1}}+\operatorname{softmax}\left(Q_{3} K_{2}^{T}\right) \overrightarrow{V_{2}}+\operatorname{softmax}\left(Q_{3} K_{3}^{T}\right) \overrightarrow{V_{3}} \
A t t_{4}(Q, K, V) & = \operatorname{softmax}\left(Q_{4} K_{1}^{T}\right) \overrightarrow{V_{1}}+\operatorname{softmax}\left(Q_{4} K_{2}^{T}\right) \overrightarrow{V_{2}}+\operatorname{softmax}\left(Q_{4} K_{3}^{T}\right) \overrightarrow{V_{3}}
\quad+\operatorname{softmax}\left(Q_{4} K_{4}^{T}\right) \overrightarrow{V_{4}}
\end{align}
$$
观察上面的式子，可以看到计算的时候，，，是被重复计算的，而且每次计算都要用到前面的和，则只与当前的有关，因此我们可以把算好的和保存下来，避免下一次自回归的时候被重复计算。这种技术我们称为KV Cache。所以，回到上一节最开始的问题，我们希望保存之前的输出避免重复运算，但实际上我们只需要保存就可以了。

PD分离

回想文本生成任务的流程，我们需要先输入prompt，让模型启动自回归，随后模型逐个输出token。由于模型每个token又要用到之前的，所以prompt的也必须计算。既然我们已经知道了prompt是什么，我们就可以利用KV Cache先把prompt的算好，之后输入到模型中，这样，模型就可以从prompt的下一个token直接自回归，加快了推理速度。

从上面的描述可以看出，模型推理的过程实际上分成了两个部分，

计算prompt的，被称为prefill（预填充）
模型自回归，被称为decode（注意与前文的decoder作区分）

这两个阶段各有特点，我们给的prompt有可能会很长，希望尽可能快的完成prefill，因此prefill明显是个计算密集型的任务。而decode阶段每次只输出一个token，计算压力并不大，所以decode阶段主要是I/O密集型任务。

由于两个任务特点不同，很容易想到可以将这两个任务放在不同的机器上运行，这种将prefill和decoder分开进行的技术就叫做PD分离。

PD的指标

PD分离主要关注两个指标，

TTFT(Time To First Token) prefill首token耗时
TPOT(Time Per Output Token) decode 每token耗时

TTFT表示prefill有多快，TPOT表示decode自回归输出token有多快。使用PD分离可以明显的提升这两个指标，因为硬件采用的是异构计算，擅长计算的机器可以专注prefill，擅长I/O的机器可以专注decode。使用混合架构时，TTFT和TPOT都会明显受损。

昇腾Ascend MindIE PD分离架构

可以看出PD分离至少需要三个node，P、D和调度器，调度器（Coordinator）是由MindIE MS实现，相当于k8s的master，P和D实例是由MindIE Server实现，相当于k8s的node。目前，仅 Atlas 800I A2 推理产品支持此特性。更多信息请参考MindIE文档PD分离服务部署

参考：彻底搞懂了Transformer解码器（图文详解） - 知乎
 大模型推理加速：看图学KV Cache - 知乎
 【超详细】【原理篇&实战篇】一文读懂Transformer-CSDN博客
 PD分离架构说明