从多少显存可以运行一个7B模型开始聊起
常常有人会聊一个问题，你如果要能够运行一个7B的模型，显卡需要多少显存才够。
感觉这个话题是很有意思的，如果你具体问题具体分析，可以先反问对方，FP8量化还是FP16量化，这时候对面当然可以反讥你，这么小的模型还需要量化？
当然我们这篇博客的重点并不在这个部分，我们先开始从传统的方式来回答这个问题，这个模型是FP16的，基础知识会告诉我们，FP16的模型，每个参数是2Byte，所以7B的模型需要14G的显存。当然，这里的显存仅仅是把模型权重加载到显存上所需要占据的显存，实际计算过程中，仍然有相当一部分需要显存占用。
这部分的内容主要包括：KV Cache、中间张量、框架开销和碎片化显存等。中间张量是可以理解的，因为计算一定会产生中间输出，中间输出在计算结束前，计算图并不会直接消失，因而结果需要空间存储。框架开销和碎片化显存也可以理解，一般一些triton算子计算的输入物理地址需要连续，所以会有这样的要求。
那么KV Cache的存在，到底是为了什么呢？其实从名字就能看出来，缓存的目的无非就是一条，用空间换时间。
那么在讨论KV Cache需要占用多少显存之前，我们需要讨论这样一个话题，为什么能用空间换时间。也就是说，KV Cache为什么能够生效。
验证KV Cache正确性
关于KV Cache，从名字就能看出来，是一个用来存储key和value的缓存，也有很多人会讲，MHA可以通过缓存KV的方式来减少实际运算，但很少有人真正讲明白，缓存的是什么KV矩阵，不需要运算的又是什么KV矩阵。这一章，我们从你的输入如何变为输出结果来看为何KV Cache能够生效，KV Cache生效又能带来多少增益。
输入《中国人能飞吗？》之后大模型会输出什么
假设，我们这里并不讨论sft之后的大模型，我们只考虑原本经过原始预训练的模型，然后这个预训练的模型样本库里有这样一句话《中国人能飞吗？当然可以！》，并且模型很好的收敛了。
假设我们模型的架构是 input->MHA->MHA->output，也就是用两个MHA来模拟网络
那么此时，我们将《中国人能飞吗？》作为输入，首先经过tokenizer，进行最初的分词以及最简单的embedding，假设模型预设的最大输入窗口是1024个token，大概这句话能够获得这样的结果，为什么mask是全1呢，因为这个batch只有一条数据，不需要对齐，当然哈，我们这里不讨论分词器，假设一个中文占据一个token。
import torch
embedding = torch.randn(1, 7, 512) # B, L, D
mask = torch.ones(1, 7, dtype=torch.long) # B, L
之后，对于你的输入的embedding，模型会首先进行prefill的运算，也就是对于这一批的数据，一次性并行计算整个Prompt序列的特征，而不是像生成阶段那样逐字递推，但是为了防止因果关系影响编码，为了确保第i个token不会偷看到第i+1个及以后的token，我们需要引入因果掩码，记为：
mask_causal = torch.triu(torch.full((L, L), float('-inf')), diagonal=1)
"""
等价于
mask_causal = torch.tensor([
[ 0., -inf, -inf, -inf, -inf, -inf, -inf],
[ 0., 0., -inf, -inf, -inf, -inf, -inf],
[ 0., 0., 0., -inf, -inf, -inf, -inf],
[ 0., 0., 0., 0., -inf, -inf, -inf],
[ 0., 0., 0., 0., 0., -inf, -inf],
[ 0., 0., 0., 0., 0., 0., -inf],
[ 0., 0., 0., 0., 0., 0., 0.]
])
"""
介绍完这些之后，我们先实现一个不带KV Cache的正经MHA，有趣的是，如果没有KV Cache，你甚至意识不到这个阶段在进行prefill，因为每一次计算都是一次prefill
import torch
import torch.nn as nn
import torch.nn.functional as F
class MultiHeadAttentionNoCache(nn.Module):
def __init__(self, hidden_dim=512, num_heads=8):
super().__init__()
self.hidden_dim = hidden_dim
self.num_heads = num_heads
self.head_dim = hidden_dim // num_heads
# 定义投影矩阵
self.W_q = nn.Linear(hidden_dim, hidden_dim, bias=False) # 偏置项为什么为0这里留到下一篇博客
self.W_k = nn.Linear(hidden_dim, hidden_dim, bias=False)
self.W_v = nn.Linear(hidden_dim, hidden_dim, bias=False)
self.W_o = nn.Linear(hidden_dim, hidden_dim, bias=False)
def forward(self,x):
B, L, D = x.shape #[1, 7, 512]
q = self.W_q(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2) # [1, 8, 7, 64]
"""
这里我们拆看看下q的计算过程
q = self.W_q(x) # [1, 7, 512]
q = q.view(B, L, self.num_heads, self.head_dim).transpose(1, 2) # [1, 8, 7, 64]
"""
k = self.W_k(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2) # [1, 8, 7, 64]
v = self.W_v(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2) # [1, 8, 7, 64]
# 计算注意力得分
scores = torch.matmul(q, k.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.head_dim)) # [1, 8, 7, 7]
"""
[1, 8, 7, 64] * [1, 8, 64, 7]/√64 = [1, 8, 7, 7]
"""
# 构造因果矩阵
mask_causal = torch.triu(torch.full((L, L), float('-inf'), device=x.device), diagonal=1) # [7, 7]
# 加上就行了，右上角会被置为负无穷，->0，(这里非常重要！！，也是KV Cache能生效的原因)
scores = scores + mask_causal
attn_weights = F.softmax(scores, dim=-1) # [1, 8, 7, 7]
# 计算输出
out = torch.matmul(attn_weights, v) # [1, 8, 7, 64]
# 拼接多头，还原输出，计算投影
out = out.transpose(1, 2).contiguous().view(B, L, D) # [1, 7, 512]
out = self.W_o(out) # [1, 7, 512]
return out
构造完成多头注意力之后，我们就构造个网络，这个网络实在太简单了，直接写出来即可，包括预测的Head
class SimpleLLMNoCache(nn.Module):
def __init__(self, hidden_dim=512, num_heads=8, vocab_size=10000):
super().__init__()
# 两层 MHA
self.layer1 = MultiHeadAttentionNoCache(hidden_dim, num_heads)
self.layer2 = MultiHeadAttentionNoCache(hidden_dim, num_heads)
# 预测下一个词的 LM Head，将 512 维度映射到 10000 词表大小
self.lm_head = nn.Linear(hidden_dim, vocab_size, bias=False)
def forward(self, x):
"""
x: [B, L, D] -> [1, L, 512]
"""
# 第一层 MHA
h1 = self.layer1(x) # [1, L, 512]
# 第二层 MHA
h2 = self.layer2(h1) # [1, L, 512]
# 输出层：映射回词表维度，计算每个位置下一个词的概率 Logits
logits = self.lm_head(h2) # [1, L, vocab_size]
return logits
虽然定义是这么定义的，但是要预测下一个词，其实logits只需要用到最后一个L的数据，在这里也就是第7个L，来进行分类任务就行了。
那么我们整体的预测下一个词的流程是什么呢？
model = SimpleLLMNoCache(hidden_dim=512, num_heads=8, vocab_size=10000)
input_prompt = torch.randn(1, 7, 512)
logits = model(input_prompt) # [1, 7, 10000]
#预测下一个词：只需要看最后一个token位置上的 logits
next_token_logits = logits[:, -1, :] # [1, 10000]
# 假设从 10000 个词里取 argmax 出来的 token 就是 “当”
next_token = torch.argmax(next_token_logits, dim=-1)
此时如果我们需要预测然呢？很简单，你已经生成当了，你可以取到当的embedding，维度也是[1, 1, 512]，将之与之前的embedding拼接，就能获得新的输入
new_input = torch.randn(1, 8, 512) # 包含 "中国人能飞吗？当"
# 再次前向传播：整条8个token的序列重算一次
logits = model(new_input) # [1, 8, 10000]
# 取第 8 个位置上的向量预测第 9 个字 "然"
next_token_logits = logits[:, -1, :] # [1, 10000]
好，此时我们关注一件事情，对于同一个初始输入，两个MHA的前几层的计算中间结果有变化吗？
import torch
import torch.nn as nn
import torch.nn.functional as F
# 1. 修正后的 MHA 模块（支持返回中间 K, V 方便观察）
class MultiHeadAttentionNoCache(nn.Module):
def __init__(self, hidden_dim=512, num_heads=8):
super().__init__()
self.hidden_dim = hidden_dim
self.num_heads = num_heads
self.head_dim = hidden_dim // num_heads
self.W_q = nn.Linear(hidden_dim, hidden_dim, bias=False)
self.W_k = nn.Linear(hidden_dim, hidden_dim, bias=False)
self.W_v = nn.Linear(hidden_dim, hidden_dim, bias=False)
self.W_o = nn.Linear(hidden_dim, hidden_dim, bias=False)
def forward(self, x):
B, L, D = x.shape
q = self.W_q(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
k = self.W_k(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
v = self.W_v(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
scores = torch.matmul(q, k.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.head_dim, dtype=x.dtype))
# 构造因果矩阵并加上
mask_causal = torch.triu(torch.full((L, L), float('-inf'), device=x.device), diagonal=1)
scores = scores + mask_causal # 修正原代码拼写错误
attn_weights = F.softmax(scores, dim=-1)
out = torch.matmul(attn_weights, v)
out = out.transpose(1, 2).contiguous().view(B, L, D)
out = self.W_o(out)
# 额外返回 k 和 v，方便观察中间变量
return out, k, v
class SimpleLLMNoCache(nn.Module):
def __init__(self, hidden_dim=512, num_heads=8, vocab_size=10000):
super().__init__()
self.layer1 = MultiHeadAttentionNoCache(hidden_dim, num_heads)
self.layer2 = MultiHeadAttentionNoCache(hidden_dim, num_heads)
self.lm_head = nn.Linear(hidden_dim, vocab_size, bias=False)
def forward(self, x):
h1, k1, v1 = self.layer1(x)
h2, k2, v2 = self.layer2(h1)
logits = self.lm_head(h2)
# 返回最终 logits 及中间过程变量
return logits, {
'layer1_k': k1, 'layer1_v': v1, 'h1': h1,
'layer2_k': k2, 'layer2_v': v2, 'h2': h2
}
# ==================== 验证实验：两次前向传播对比 ====================
# 0. 固定随机种子，初始化模型
torch.manual_seed(42)
model = SimpleLLMNoCache(hidden_dim=512, num_heads=8, vocab_size=10000).eval()
# 模拟前 7 个 Token（"中国人能飞吗？"）的 Embedding 矩阵
prompt_7_tokens = torch.randn(1, 7, 512)
# 模拟刚预测出来的第 8 个 Token（"当"）的 Embedding
token_8_embedding = torch.randn(1, 1, 512)
# 拼接成 8 个 Token（"中国人能飞吗？当"）
prompt_8_tokens = torch.cat([prompt_7_tokens, token_8_embedding], dim=1)
# ------------------------------------------------------------------
# 第一次执行：输入 7 个 Token
# ------------------------------------------------------------------
with torch.no_grad():
logits_pass1, inter_pass1 = model(prompt_7_tokens)
# ------------------------------------------------------------------
# 第二次执行：输入 8 个 Token
# ------------------------------------------------------------------
with torch.no_grad():
logits_pass2, inter_pass2 = model(prompt_8_tokens)
# ------------------------------------------------------------------
# 打印对比结果（截取第二次执行中前 7 个 Token 的中间输出进行对比）
# ------------------------------------------------------------------
print("=" * 65)
print("对比两次执行中【前 7 个 Token】对应位置的中间计算结果：")
print("=" * 65)
def compare_tensors(name, tensor1, tensor2_slice):
# 计算最大绝对误差
max_diff = torch.max(torch.abs(tensor1 - tensor2_slice)).item()
is_equal = torch.allclose(tensor1, tensor2_slice, atol=1e-7)
print(f"[{name}]")
print(f" └─ Shape 对比: Pass1 {list(tensor1.shape)} vs Pass2[:7] {list(tensor2_slice.shape)}")
print(f" └─ 最大绝对误差 (Max Diff): {max_diff:.8e}")
print(f" └─ 数值完全一致: {is_equal}\n")
# Layer 1 的 K 矩阵 (仅对比前 7 个位置)
compare_tensors("Layer 1 - K 矩阵", inter_pass1['layer1_k'], inter_pass2['layer1_k'][:, :, :7, :])
# Layer 1 的 V 矩阵 (仅对比前 7 个位置)
compare_tensors("Layer 1 - V 矩阵", inter_pass1['layer1_v'], inter_pass2['layer1_v'][:, :, :7, :])
# Layer 1 的输出状态 h1 (仅对比前 7 个位置)
compare_tensors("Layer 1 - 输出隐藏状态 h1", inter_pass1['h1'], inter_pass2['h1'][:, :7, :])
# Layer 2 的输出状态 h2 (仅对比前 7 个位置)
compare_tensors("Layer 2 - 输出隐藏状态 h2", inter_pass1['h2'], inter_pass2['h2'][:, :7, :])
# 最终 Logits (仅对比前 7 个位置)
compare_tensors("LM Head - 最终 Logits", logits_pass1, logits_pass2[:, :7, :])
最终执行的结果是
=================================================================
对比两次执行中【前 7 个 Token】对应位置的中间计算结果：
=================================================================
[Layer 1 - K 矩阵]
└─ Shape 对比: Pass1 [1, 8, 7, 64] vs Pass2[:7] [1, 8, 7, 64]
└─ 最大绝对误差 (Max Diff): 0.00000000e+00
└─ 数值完全一致: True
[Layer 1 - V 矩阵]
└─ Shape 对比: Pass1 [1, 8, 7, 64] vs Pass2[:7] [1, 8, 7, 64]
└─ 最大绝对误差 (Max Diff): 0.00000000e+00
└─ 数值完全一致: True
[Layer 1 - 输出隐藏状态 h1]
└─ Shape 对比: Pass1 [1, 7, 512] vs Pass2[:7] [1, 7, 512]
└─ 最大绝对误差 (Max Diff): 2.38418579e-07
└─ 数值完全一致: True
[Layer 2 - 输出隐藏状态 h2]
└─ Shape 对比: Pass1 [1, 7, 512] vs Pass2[:7] [1, 7, 512]
└─ 最大绝对误差 (Max Diff): 1.49011612e-07
└─ 数值完全一致: True
[LM Head - 最终 Logits]
└─ Shape 对比: Pass1 [1, 7, 10000] vs Pass2[:7] [1, 7, 10000]
└─ 最大绝对误差 (Max Diff): 1.26659870e-07
└─ 数值完全一致: True
你会发现，KV的结果是完全一致，如果不考虑浮点计算精度导致的问题，logic的结果也是完全一致的，这是为何呢！
因为因果掩码的存在，第二次计算中前7个token的计算output完全不受到第8个token的影响，也就是说，对于第二次的计算MHA的时候，前七个token的embedding是不会变的。这里我让ai画了个图，我觉得是很好理解的。
首先是为什么KV能够缓存：
第一次前向传播 (Pass 1 - 输入 7 个 Token):
[ X_1 ] [ K_1 ] [ V_1 ]
[ X_2 ] [ K_2 ] [ V_2 ]
[ ... ] * W_k, W_v = [ ... ] 和 [ ... ]
[ X_7 ] [ K_7 ] [ V_7 ]
---------------------------------------------------------------------
第二次前向传播 (Pass 2 - 输入 8 个 Token):
[ X_1 ] [ K_1 ] (完全一致) [ V_1 ] (完全一致)
[ X_2 ] [ K_2 ] [ V_2 ]
[ ... ] * W_k, W_v = [ ... ] 和 [ ... ]
[ X_7 ] [ K_7 ] [ V_7 ]
=====================================================================
[ X_8 ] <-- 新增 [ K_8 ] <-- 新增 [ V_8 ] <-- 新增
其次是为什么这个MHA的输出两次也是相同的，也就是进行了Attention以及过了后面的W_o依然相同
这两次计算，既然K，V是相同的，唯一可能影响结果的只有Query了，但是由于因果掩码的存在
Key 1 Key 2 ... Key 7 | Key 8 (新增)
+-------------------------------+----------------
Query 1 | S_11 -inf ... -inf | -inf <-- 被 Mask 挡住！
Query 2 | S_21 S_22 ... -inf | -inf <-- 被 Mask 挡住！
... | ... ... ... ... | -inf
Query 7 | S_71 S_72 ... S_77 | -inf <-- 被 Mask 挡住！
-------------+-------------------------------+----------------
Query 8(新增)| S_81 S_82 ... S_87 | S_88
我们可以发现 （Q·K_T）·V之前，进行了掩码操作，Q_8以及K_8的结果，无法影响到计算O_0到O_7的最终计算过程，所以此时，说明，对于下一个MHA，输入的前7个编码是一定一致的，既然输入一致，那么下个MHA的中间结果也是一直的，也就是说不仅是对于这一个MHA可以缓存KV，后面的任何MHA的KV都是可以缓存的！
至此，我们已经成功论证了KV Cache的有效性！
放飞思想，学习理当从愉快的阅读开始