LLama2中，流程中数据大小的变换如下
Transformer模块
第一次输入，进行prefill，输入x维度为[1, 8, 4096]
1. 构建wq,wk,wv,wo,尺寸均为[4096,4096]， 与x点乘，得到xq, xk, xv
2. 构建KV cache， 尺寸为 [batch size, max_seq_len, local_kv_heads, head_dim],对应 [1, 8, 32, 128]
3.基于kv cache构造 keys, alues,对应的尺寸还是[1,8,32,128]
4. 在最后两个维度对于xq和key进行点乘，得到scores，维度变成【1， 32， 8， 8】
5. 将mask与scores相加
6. 对于scores进行softmax
7. 将scores [1, 32, 8, 8]与values [1, 32, 8, 128]进行乘法
8. 得到output [1, 8, 4096]
9. 将output再与wo进行乘法[1, 8, 4096]
10. 接下来对于输出进行 ffn_norm的操作
Feedforward模块
11.然后进行feed_forward.得到当前transformer模块的输出 [1, 8, 4096]
feed_forward的操作如下，虽然代码很小，但是计算量却很大。
def forward(self, x):
return self.w2(F.silu(self.w1(x)) * self.w3(x))其中，w1的维度为[11008, 4096], w2的维度为[4096, 11008], w3的维度为[11008, 4096]
kv cache的表达如下
self.cache_k = torch.zeros(
(
args.max_batch_size,
args.max_seq_len,
self.n_local_kv_heads,
self.head_dim,
)
).cuda()
self.cache_v = torch.zeros(
(
args.max_batch_size,
args.max_seq_len,
self.n_local_kv_heads,
self.head_dim,
)关于kv cache的细节讨论
llama2设定 local_kv_heads为32，head_dim为128。所以，kv cache的尺寸为 [1, 512,32, 128] * 2
对于一个batch的数据来说哦，因为llama2 7B 包含32个transformer，所以，当使用FP32表达时， 对应一个batch的kv cache的大小为128 * 32 * 128 *2 * 32 * 4byte= 0.5GB.
这里，也可以看到几个变量：
* 当batch变大时，kv cache线性增长
* 当batch 的最大长度增大时， Kv cache线性增长。
参考链接：