大模型推理显存和计算量估计方法

2025年6月23日
| 没有评论

最近做吞吐量调试涉及到输入batch_size的设置，为了把算力和显存用起来，同时不触发out of memory，需要提前估计大模型推理过程中的显存占用，我参考了MindIE官网的这个文档完成了估计：https://www.hiascend.com/document/detail/zh/mindie/20RC2/mindieservice/servicedev/mindie_service0105.html 。

显存估计

大模型推理的过程中，显存主要用来存储kvcache。kvcache的大小和token的数量成正比，我们首先来看一下单个token的kvcache怎么计算。假设单token输入经过embedding层之后的维度是hidden_size，那么接下来就要和k权重矩阵和v权重矩阵相乘，kv的权重矩阵shape是[hidden_size, hidden_size]，所以计算得到的k和v的维度是hidden_size。需要注意的是，虽然这里kv矩阵只有1个，但实际上现在大部分模型都采用了多头注意力，注意力头的数量是num_attention_heads，单个注意力头的权重矩阵是[hidden_size， hidden_size/num_attention_heads]。此外，如果采用分组kv头，也就是模型的config.json文件中包含num_key_value_heads参数，那么kv的权重矩阵shape就是[hidden_size, (hidden_size/num_attention_heads)*num_key_value_heads]，这样的话要缓存的kv的维度就是(hidden_size/num_attention_heads)*num_key_value_heads。由于大模型一般包含多层transformer，所以还需要乘以层数。总的来说，单个token占用的kvcache显存大小计算公式如下： 

import numpy as np
# take qwen2.5-7B as example
hidden_size = 3584
num_attention_heads = 28
num_hidden_layers = 28
num_key_value_heads = 4  # if not grouped kv cache, num_key_value_heads equals num_attention_heads
kvcache_dtype_byte = 2  # bf16/fp16 -> 2, int8 -> 1
cache_num = 2  # k cache and v cache

one_token_cache = hidden_size / num_attention_heads * num_key_value_heads * num_hidden_layers * kvcache_dtype_byte * cache_num
print(f"one_token_cache is: {one_token_cache} byte.")

计算了单个token的kvcache显存大小后，我们就可以计算整个序列所占的显存大小了： 

input_length = 1024
output_length = 1024
batch_size = 16
allocate_mem = one_token_cache * (input_length + output_length) * batch_size / (1024*1024*1024) # GB
print(f"allocate memory should be larger than {allocate_mem} G.")

还可以估计所能支持的最大batch_size： 

# estimate max batch_size for qwen2.5-7B
model_size = 7  # B
total_mem = 64  # 64G for each npu
mem_coefficient = 0.8
mem_for_kvcache = (total_mem - model_size*kvcache_dtype_byte)*mem_coefficient*1024*1024*1024  # byte
kvcache_block_size = 128
max_block_num = np.floor(mem_for_kvcache/(kvcache_block_size*one_token_cache))
print("max_block_num: ", max_block_num)
one_sequence_block_num = np.ceil((input_length + output_length) / kvcache_block_size)
max_batch_Size = np.floor(max_block_num / one_sequence_block_num)
print("max_batch_Size: ", max_batch_Size)

在估计最大batch_size的时候，我们先根据显卡的最大容量（64G）和模型大小估计了可用于缓存kvcahe的空间mem_for_kvcache，然后估算了可以分配多少个kvcache_block。接着计算一个sequence所需要的kvcache_block数量，最后用总的kvcache_block数量除以一个sequence所需要的kvcache_block数量，得到的就是支持的最大batch数量。

上面的代码运行结果如下： 

one_token_cache is: 57344.0 byte.
allocate memory should be larger than 1.75 G.
max_block_num:  5851.0
max_batch_Size:  365.0

为了验证理论分析是否正确，我们用MindIE跑一下qwen2.5-7B模型： 

bash run.sh pa_bf16 performance [[1024,1024],[256,256]] 16 qwen /home/jinxiulang/qwen2.5/Qwen2.5_7B_Instruct 1

运行日志截图如下：

可以看到，日志中包含了“kv cache will allocate 1.75GB memory”，和我们上面估算的是一致的。

计算量估计

为了估算推理执行耗时，还需要估计模型推理消耗的计算量，然后结合芯片算力估计时延。

Transformer模型的计算量主要来自自注意力机制、前馈网络（FFN）和最后的lm_head层。假设模型参数如下：

  • L：Transformer层数。
  • H：隐藏层大小（隐藏维度）。
  • I：FFN中间层大小（通常I>H）。
  • V：词表大小。
  • S：prefill输入序列长度。
  • T：序列总长度（prefill+生成token）。

Prefill计算量

计算每层Transformer的FLOPs：

首先计算自注意力层。第一步是求q、k、v，计算公式是input*weight，input的shape是[S, H]，weight的shape是[H, H]，所以q/k/v要做SH次向量内积，每次向量内积要做H次乘法和（H-1）次加法，近似于2H，所以q、k、v所有的计算量是3*SH*2H=6SH^2。需要注意的是，如果采用的是分组kvcache，那么计算q、k、v的时候，H要换成(H/num_heads)*num_kv_heads，但是在估算的时候可以近似为H；然后是Q*K（实际上是Q乘以K的转置），输入shape是[S, H]和[H, S]，计算量是2S^2H；接着是Q*K*V，输入shape是[S, S]和[S, H]，计算量是2S^2H；最后是输出投影，输入shape是[S, H]和[H, H]，计算量是2SH2；所以自注意力层的计算量是8SH2+4S^2H。

然后计算FFN层的计算量，FFN层包含一个升维层、一个gate层和一个降维层，计算量分别为2SHI、2SHI、2SIH，所以总计算量为6SHI。

所以每个transformer层的计算量为8SH2+4S2H+6SHI。再加上最后的lm_head，prefill的计算量为L(8SH2+4S2H+6SHI)+2SHV。

Decode计算量

和prefill相比，decode的主要变化是输入序列长度为1。

首先计算自注意力层。计算qkv的计算量是6H^2。计算q*k的时候，由于要把缓存的k也加上，假设当前序列长度是T，那么输入shape是[1, H]和[H, T]，所以计算量是2HT；接着计算qkv，输入shape是[1,T]和[T, H]，计算量是2HT。所以自注意力层的计算量是8H^2+4HT。

然后计算FFN层的计算量，参考prefill的过程，可知是6HI。

所以每个transformer层的计算量为8H2+4HT+6HI。再加上最后的lm_head，decode的计算量为L(8H2+4HT+6HI)+2HV。需要注意的是，这里面的T是变化的，如果要计算K个decode过程的平均耗时，可以取T=(S+S+K)/2进行估计。

实验验证

为了验证我们的理论公式，我们继续基于qwen2.5-7B进行验证，还是运行： 

bash run.sh pa_bf16 performance [[1024,1024],[256,256]] 16 qwen /home/jinxiulang/qwen2.5/Qwen2.5_7B_Instruct 1

输出结果部分截图如下：

在输入/输出取1024/1024，batch_size=16的情况下，首token时延1088.62ms，decode平均时延18.77ms。

我们把理论计算公式用python实现： 

# estimate calculation amount for qwen2.5-7B
# prefill
L = num_hidden_layers
S = input_length
H= hidden_size
I = 18944
V = 152064
attention_cal_prefill = L*(8*S*H*H + 4*S*S*H)
forward_cal_prefill = L*6*S*H*I
head_cal_prefill = 2*S*H*V
total_cal_prefill = (attention_cal_prefill + forward_cal_prefill + head_cal_prefill)*batch_size
print(f"Total calculation amount for prefill is {total_cal_prefill:.3e}, the percentages of attention layer,"
      f" forward layer and head layer are: {attention_cal_prefill/total_cal_prefill}, {forward_cal_prefill/total_cal_prefill}, {head_cal_prefill/total_cal_prefill}")

# decode
T = input_length + output_length / 2
attention_cal_decode = L*(8*H*H + 4*H*T)
forward_cal_decode = L*6*H*I
head_cal_decode = 2*H*V
total_cal_decode = (attention_cal_decode + forward_cal_decode + head_cal_decode)*batch_size
print(f"Total calculation amount for averaged decode is {total_cal_decode:.3e}, the percentages of attention layer,"
      f" forward layer and head layer are: {attention_cal_decode/total_cal_decode}, {forward_cal_decode/total_cal_decode}, {head_cal_decode/total_cal_decode}")

npu_cal_ability = 300*1e+12

print(f"The calculation time for prefill and decode are: {(total_cal_prefill / npu_cal_ability)}, {total_cal_decode / npu_cal_ability}")

注意，代码中的npu_cal_ability是我们使用的npu单卡算力，300T Flops左右。

输出结果如下： 

Total calculation amount for prefill is 2.586e+14, the percentages of attention layer, forward layer and head layer are: 0.20832364566379913, 0.7226226458963032, 0.06905370843989769
Total calculation amount for averaged decode is 2.558e+11, the percentages of attention layer, forward layer and head layer are: 0.21849896717925177, 0.7133348634381456, 0.06816616938260271
The calculation time for prefill and decode are: 0.8620572025378134, 0.000852813851306666

可以看到，prefill的预估时间为862ms，和实测的1088.62ms比较接近，但是增量平均时延为0.85ms，远小于实测的18.77ms。主要有以下原因：

1，AI芯片的算力参数为峰值算力，一般情况下AI芯片的算力利用率在60%左右；

2，AI芯片的算力大部分来自矩阵乘法单元，全量计算都是大矩阵运算GEMM，可以充分利用AI Core的能力，但是增量计算都是小矩阵运算（特别是batch_size=1的时候，退化为向量矩阵运算GEMV），导致算力利用率很低；

3，token时延除了包含计算时间，还有内存搬运时间、软件栈之间的数据传输时间等等，对于decode这种运算时间短的场景，其他环节的时延会占很大的比例。

此外，对于分布式推理场景，不能仅凭计算量来估计时延，因为通信算子的耗时往往会较大，比如MOE结构模型中的alltoall通信在推理过程中可能占总时延的30%。

本文由博客一文多发平台 OpenWrite 发布！

flask, python, tornado, 人工智能AI