一、推理大致流程总结
首先我们肯定要输入prompt，prompt经过tokenizer转换成token，这里的token是经过tokenizer的词汇表映射的，而且token是能够被模型理解的。然后在gpu的支持下llm会逐步生成输出，每次生成一个token后，这个token会作为输入再次输入到llm中。
在llm开始生成输出之前，会对输入的propmt进行处理和缓存，其实这里有个首token延迟的问题，后续在详细说明这个问题。
在生成token的过程中模型会依赖已经生成的内容和cache逐步生成新的token，这也就是kvcache，以空间换时间加快生成速度。
最后生成的token会通过detokenizer会被转回人类可读的文本形成输出让我们查看。
推理流程：
推理有两个阶段，第一个是prefill阶段也就是输入prompt处理的阶段，会生成cache
第二个是decode后续生成新token的阶段，会利用prefill的cache以及阶段本身产生的cache
那么在prefill阶段会发生什么？
输入的prompt会经过分词器分解为token，每个token都有唯一的tokenid，然后输入的tokenid会经过嵌入层（embedding layer）转为一个高维连续的向量，比如说768维或者1024维向量。
具体的代码如下：

from torch import nn
from transformers import AutoConfig
from transformers import AutoTokenizer

model_ckpt = “Bert-base-uncased” #分词器模型
Tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
Text = “time flies like an arrow”

inputs = tokenizer(text, return_tensors="pt", add_special_tokens=False) #映射为tokenid
config = AutoConfig.from_pretrained(model_ckpt)

token_emb = nn.embedding(config.vocab_size,config.hidden_size) #这里配置embedding层的词表和维度
inputs_embeds = token_emb(inputs.input_ids) #这里按照配置好的config来进行token转embedding的转换

转成embedding（词向量）的过程是通过torch.embedding函数转的。当然在prompt文本转token过程中还会有些分隔词（也就是add_special_token)如果是false的话就去除了分词结果中的 [CLS] 和 [SEP]，true的话token中还有特殊的token
然后就是attention的计算了，这里就是transformer的内容，包括一系列attention计算过程，然后加位置编码等。
这里的中间结果会被保存下来，缓存的内容包含已经计算好的embedding向量和注意力全中等信息。
所以总结下来信息是：
（1）输入文本Token化，转成embedding
（2）生成查询（Query）、键（Key）、值（Value）向量：这种生成方式是通过矩阵乘法实现的（例如：嵌入矩阵 * W_Query = Query向量），每个token的查询、键、值向量都会作为模型计算注意力的基础。
这个W_Query，W_Key,W_Value几个矩阵都是可训练参数，随着模型训练会不断调整，以便模型能够更准确地理解输入序列中的关系。所以多头注意力机制一般会有不同的W_query,W_key,W_value。
（3）计算注意力分数：在注意力机制中，模型会计算每个查询向量和其他token的键向量之间的点积（或相似度）。这个点积结果表示每个token对其他token的“注意力分数”。
（4）权重化和加权求和：注意力分数会经过Softmax处理，将分数转换为概率分布，然后用这些分数对相应的值向量进行加权。
（5）多头注意力机制：为了提高模型的表达能力，LLM通常会采用多头注意力。每个头代表一组独立的Query、Key和Value矩阵，能够关注输入中的不同特征。
通过多个注意力头模型可以从不同角度理解输入序列中的信息，从而捕捉更丰富的上下文信息。
（6）生成阶段将attention等信息传到模型中，基于这些信息生成新的token，这是个自回归的过程，然后逐步生成token直到完整的句子。
（7）gpu加速和缓存，这其中注意力计算会占用很大的计算资源，尤其在多头注意力和长序列输入过程中，因此llm会利用gpu进行并行计算，在gpu内存中缓存query，key和value等矩阵提高计算效率。所以如果没有kv cache这个机制的话每生成一个新的token都需要重新计算当前序列中所有token的key-value矩阵，由于llm推理是个自回归过程，所以会不断重复计算浪费计算资源。
归根结底从输入prompt到llm生成token其实只有两个阶段prefill和decode/generation阶段。一个请求过来后实际的执行顺序就是Prefill+decode。但实际过程中肯定不只有一个请求打过来，如果当第一个请求执行到decode阶段的时候，如果第二个请求打过来，这个时候肯定不能拒接第二个请求，而是同样第二个请求也开始执行，此时第二个请求在prefill阶段而第一个请求在decode阶段，也就是下面这幅图：

但是这里需要注意prefill和decode两个阶段各自的运行特性和资源需求有显著的不同：
Prefill阶段：
● Prefill阶段会并行处理输入的所有token，这种处理方式使得即使在较小的batch size下也能打满GPU的利用率
● 由于在prefill阶段需要处理长输入（如512 tokens），所以这个阶段的计算开销很大，显卡利用率很容易打满了
● 增大batch size时，prefill阶段每个token的处理开销几乎保持不变，这意味着prefill的效率在小batch size时就已经非常高，说明开销是固定的
Decode 阶段：
● Decode阶段是自回归的，每次只生成一个token，因此这一阶段的GPU利用率通常较低
● IO密集型：Decode过程中需要频繁地读取KV Cache，导致IO开销较大。即使输入的长度始终为1，反复的KV Cache访问也使得这一阶段成为IO密集型
● 扩大batchsize可以显著降低decode阶段的开销，因为更大的batchsize能更有效地分摊固定的IO读写成本，不过开再大也不能完全打满GPU，毕竟KV Cache的读写开销过大，导致decode阶段难以成为计算密集型。

二、推理瓶颈/细节

1、flash attention
然后还有flash attention，为啥要说flash attention呢，因为现在的llm基本上架构都差不多，（其实Qwen2-audio就用了flash attention）都是前馈，激活层，BN层，自注意力这些东西。其中自注意力这块很重要，它是让模型能够理解输入的token上下文关系，但是self attention会随着输入token的数量占据非常大的gpu内存。
$Attn(X)=V\ast Softmax（Q{K}^T)$，这里的QKV都是W_q，W_k，W_v与输入X的乘。
又因为llm有多个注意力头，所以会并行的执行多个自注意力计算，如果输入的序列越来越长那么存储$Q{K}^T$的内存就会越来越大。
这里有一种新的方法来避开QK^T矩阵

可以看到传统的自注意力计算需要从HBM中加载qk，在把S=QK矩阵乘写到HBM中。从HBM读S，计算S的Softmax=P，写到HBM中。从HBM中加载P和V，计算PV矩阵乘，写到HBM中。这里的IO开销很大。
这套算法的HBM访存时间复杂度O(Nd+N^2) 在gpt2中N=1024,d=64.hbm带宽比较小，所以访存成为瓶颈

flash attention的主要思想就是减少HBM的访问，将QKV切分为小块后放入SRAM进行计算,原始的softmax数值不稳定，所以flash attention采用safesoftmax。
标准的self attention需要计算$softmax(Q{K}^T)\ast V$,其中 Q、K 和 V 分别是查询、键和值的投影矩阵。假设输入的长度为 N，那么$Softmax(Q{K}^T)$是一个 $N\times N$ 的矩阵。随着 N 增加，计算和存储 $Q{K}^T$所需的内存将成二次方增长$O(N^2)）$，这是长序列推理中的主要瓶颈。flash attention将$Q{K}^T$的计算过程分解为多个小块，避免存储完整的$Q{K}^T$矩阵减小内存消耗。

由上图可以看到计算小块之间的$Q_i{K}_i^T$
,而不是整个的$Q{K}^T$，这样的话不需要存储$N\ast N$的矩阵，只需要存一个小块的结果就可以。
存完小块的结果还需要softmax计算，对于每个小块结果我们可以直接执行softmax，得到这块的结果$O_i$。当然论文中的细节还是很多，例如safesoftmax如何找最大值，如何计算，但是推理大致流程不需要介绍这么细节。
然后为每个小块维护归一化统计量，例如$s_{ij}^a$
和$s_{ij}^b$。它们用于在不同块之间协调计算，使得最终的输出保持一致。这些统计量在每次迭代中都需要重新计算，以确保不同小块之间的计算一致性。
最后计算当前块的$attention：O_i\leftarrow s_{ij}^a\ast O_i+s_{ij}^b\ast V_j\times Softmax(Q{K}_{i,j}^T)$
在第6行可以看到K，V还有Q等参数从HBM转到SRAM上了，后续计算也是从SRAM上计算。虽然Flash Attention通过分块计算看似更复杂，但由于追踪并归一化Softmax统计量s_{ij}^a
和s_{ij}^b，最终结果在数值上与标准的自注意力计算保持一致。因此，Flash Attention提供了数学上相同的输出，且仅需线性增长的内存，而非正常指数级。
flash attention2（https://zhuanlan.zhihu.com/p/645376942）也是类似思路，但是更多的是从cuda和GPU内部架构方面进行优化，所以此处不在提及。
2、chunked prefill
刚才在推理流程第一节部分也说过，Prefill效率高但性能一般固定，在小batch size下即可达到高GPU利用率，对batch size不敏感。但Decode效率低且受batch size影响大，也就是说Decode阶段需要大batch size才能提高GPU利用率，但同时受限于KV Cache的性能瓶颈。此时一个新技术就出场
待填坑（https://arxiv.org/pdf/2308.16369）
3、kv cache
我们知道llm推理是个自回归过程，所以模型会迭代地输入一个序列，采样下一个token，将该token附加到输入序列中，并继续此过程，直到LLM生成一个表示生成结束的token。所以tokens不会依赖未来的tokens，那么映射到Q,K,V来说，q_i向量（第i个token的查询向量）不会与任何j>i的k_j,v_j相关联。相反，q_i仅仅与之前的键值向量k_m,v_m ,其中m∈{0,1,2…i-1}，所以为了减少计算次数我们可以缓存之前的键值向量。
我们将通过在每次前向传递时检索并转发键值缓存来告诉LLM使用键值缓存。在Transformers库中，可以通过将 use_cache 参数传递给前向调用来检索键值缓存，然后将其与当前token一起传递。

past_key_value = None
generated_tokens = []
next_token_id = tokenizer(prompt,return_tensors="pt")["input_ids"].to("cuda")

for _ in range(5):
  next_logits,past_key_values = model(next_token_id,past_key_value=past_key_values,use_cache=True).to_tuple()
  next_logits = next_logits[:, -1:]
  next_token_id = torch.argmax(next_logits, dim=-1)

  print("shape of input_ids", next_token_id.shape)
  print("length of key-value cache", len(past_key_values[0][0]))  # past_key_values的形状为[num_layers, 0 表示 k, 1 表示 v, batch_size, length, hidden_dim]
  generated_tokens.append(next_token_id.item())

输出如下，可以看到传入的仅仅为当前token而非全部token，但是key-value cache长度却始终+1.

shape of input_ids torch.Size([1, 1])
length of key-value cache 20
shape of input_ids torch.Size([1, 1])
length of key-value cache 21
shape of input_ids torch.Size([1, 1])
length of key-value cache 22
shape of input_ids torch.Size([1, 1])
length of key-value cache 23
shape of input_ids torch.Size([1, 1])
length of key-value cache 24

所以使用kv cache意味从原来计算$Q{K}^T转换为计算q_c{K}^T$，计算量大大减小提升了推理速度。而且最大内存需求也并不随着token的数量呈现二次增长，而是呈线性增长。
但是有时候使用kvcache也会在推理时出现问题（例如llama：https://github.com/huggingface/transformers/issues/25420#issuecomment-1775317535）
在多轮对话中，kvcache的作用更为明显。
User: How many people live in France?
Assistant: Roughly 75 million people live in France
User: And how many are in Germany?
Assistant: Germany has ca. 81 million inhabitants
这里是两轮对话，llm进行了两次decode。其中第一次输入时kvcache为空，输入prompt为“How many people live in France?”，模型通过自回归生成“Roughly 75 million people live in France”。解码过程中增加kvcache。
第二次问答中，输入prompt为And how many are in Germany?，这时候其实对于模型来说输入的是：User: How many people live in France?
Assistant: Roughly 75 million people live in France
User: And how many are in Germany?
这三句中，因为前两句的kv已经被计算过，所以有效的输入只有最后一句。经过一系列处理，新计算的kv值会拼接到前一次的kvcache后面。但是需要注意，长对话的过程中kvcache也是非常占用内存的！
下面通过矩阵的形式来说明kvcache：
首先假设目前只处理一个长度为t的序列（batchsize=1）
● 在整个过程中，输入序列中的每个token都由一个稠密向量表示。
● 注意力层的输入是一系列稠密向量，每个输入token对应一个，由前一解码器块生成。
● 对于每个输入向量，注意力层生成一个相同维度的输出稠密向量。
先来看看不使用cache的过程，假设模型生成遥遥领先四个字：
当模型生成第一个“遥”字时，input=“”, ""是起始字符。Attention的计算如下：

当模型生成第二个“遥”字时，input=“遥”, Attention的计算如下：

可以看到最后结果中第一行attention被重复计算。
所以第一行的Att_1(Q,K,V） = softmax(Q_1K_1^{T)V_1是第一次attention计算的值，会发现Q_1K_2}T这个值会被mask掉。
● 所以Q1在第二步参与的计算与第一步是一样的，而且第二步的V1也仅仅依赖于 Q1 ，与 Q2 毫无关系。
● V2 的计算也仅仅依赖于 Q2 ，与 Q1 毫无关系。
当模型生成第三个“领”字时，input="遥遥"Attention的计算如下，结果依次类似于step2的结果

所以得出规律，Att只有Q有关。

当模型生成第四个“先”字时，input="遥遥领"Attention的计算如下：

公式也和之前类似，不再赘述。由此我们可以得出结论：

1. 当前计算方式存在大量冗余计算。
2. Att_k只与 Q_k 有关。
3. 推理第 x_k 个字符的时候只需要输入字符 x_{k-1}即可。
   如果使用cache的话，过程会是什么样？
   
   
   
   以此类推…最后只需要concat起来就是完整的结果。

推理summary

推理有两个阶段，第一个是prefill阶段也就是输入prompt处理的阶段，会生成cache
第二个是decode后续生成新token的阶段，会利用prefill的cache以及阶段本身产生的cache

from torch import nn
from transformers import AutoConfig, AutoTokenizer

model_ckpt = "bert-base-uncased"  # 分词器模型
tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
text = "time flies like an arrow"

inputs = tokenizer(text, return_tensors="pt", add_special_tokens=False)  # 映射为 token ID
config = AutoConfig.from_pretrained(model_ckpt)

token_emb = nn.Embedding(config.vocab_size, config.hidden_size)  # 配置 embedding 层的词表大小和维度
inputs_embeds = token_emb(inputs.input_ids)  # 根据配置好的 config，将 token ID 转换为 embedding 向量

参考
https://huggingface.co/docs/transformers/en/perf_infer_gpu_one#flashattention-2
https://arxiv.org/pdf/2205.14135
https://zhuanlan.zhihu.com/p/676655352
https://huggingface.co/docs/transformers/main/en/llm_tutorial_optimization#32-the-key-value-cache
https://blog.csdn.net/KIDGIN7439/article/details/131293535
https://huggingface.co/docs/transformers/llm_tutorial#generate-text
https://www.youtube.com/watch?v=z2M8gKGYws4&list=PL_lsbAsL_o2B_znuvm-pDtV_cRhpqZb8l&index=24