前面两章，我们探讨了怎样控制模型生成：采样策略和停止条件。接下来，研究一下怎样提升模型推理的性能。
GPT 模型自回归生成有一个性能瓶颈：每生成一个新词元，都要重新计算所有词元的注意力。假设已经生成了 100 个词元，要生成第 101 个词元就要计算 101 个词元的注意力，生成每个词元的时间复杂度为 。
键值缓存¶
键值缓存（KV Cache）的核心思想：把已生成词元的键（Key）和值（Value）缓存起来。生成新词元时只计算查询（Query），再和缓存的键和值一起做注意力。
实现原理¶
自注意力中，每个词元的键和值只取决于它自身及之前的词元，与后续词元无关。因此：
已生成词元的键和值可以缓存，不需要重复计算；
新词元只需要计算自己的查询，和缓存的键和值一起做注意力；
缓存随着生成过程不断增加，但是每个词元只计算一次。
实现要点¶
缓存结构：每个 Transformer 模块存一份键和值，其中键已转置，格式为
(批大小, 头数, 每头嵌入尺寸, 已生成长度)；值的格式为(批大小, 头数, 已生成长度, 每头嵌入尺寸)；增量计算：新词元的键和值和缓存拼接，再计算注意力；
位置编码：新词元的位置编码要加上偏移量（已生成长度）。
import time
import numpy as np
from src.core import *
from src.gpt import *np.random.seed(42)class GPTEmbeddingWithCache(GPTEmbedding):
def forward(self, x: Tensor, offset=0):
token = self.embedding(x)
# 位置偏移
position = self.positional_embedding(Tensor(range(offset, offset + x.shape[1])))
return self.dropout(token + position)GPT 缓存注意力层¶
GPT 缓存注意力层（GPTAttentionWithCache）是键值缓存的核心：
无缓存：输入值是当前位置之前的整个上下文尺寸，正常计算；
有缓存：输入值只有当前位置的新词元，键和值需要和缓存拼接，再计算注意力。
有缓存时不要因果掩码，因为新词元只与已生成词元计算注意力，天然满足因果关系。
class GPTAttentionWithCache(GPTAttention):
def forward(self, x: Tensor, use_cache=False, cache=None):
norm = self.normalize(x)
head_embedding_size = self.embedding_size // self.heads
multi_head_shape = (-1, x.shape[1], self.heads, head_embedding_size)
shape = (-1, x.shape[1], self.embedding_size)
query = self.query(norm).reshape(multi_head_shape).transpose((0, 2, 1, 3))
key = self.key(norm).reshape(multi_head_shape).transpose((0, 2, 3, 1))
value = self.value(norm).reshape(multi_head_shape).transpose((0, 2, 1, 3))
# 有缓存时，输入值只有最新词元，需要和缓存拼接
if use_cache and cache is not None:
key = Tensor(np.concatenate([cache[0].data, key.data], axis=3))
value = Tensor(np.concatenate([cache[1].data, value.data], axis=2))
# 更新缓存
new_cache = (key, value) if use_cache else None
scale = Tensor(np.array(1.0 / np.sqrt(head_embedding_size)))
scores = query @ key * scale
# 无缓存时需要因果掩码
if cache is None:
scores = self.mask(scores)
weights = self.softmax(scores)
output = self.output((weights @ value).transpose((0, 2, 1, 3)).reshape(shape))
return x + self.dropout(output), new_cache
GPT 缓存 Transformer 模块¶
GPT 缓存 Transformer 模块（GPTTransformerWithCache）把缓存作为参数传递给注意力层，并返回注意力层产生的新缓存。
class GPTTransformerWithCache(Composite):
def __init__(self, embedding_size, heads):
self.attention = GPTAttentionWithCache(embedding_size, heads)
self.feed_forward = GPTFeedForward(embedding_size)
super().__init__([self.attention,
self.feed_forward])
# 推理函数，传递缓存
def forward(self, x: Tensor, use_cache=False, cache=None):
x, new_cache = self.attention(x, use_cache, cache)
return self.feed_forward(x), new_cacheclass GPTWithCache(Composite):
def __init__(self, vocab_size, context_size, embedding_size, heads, blocks):
self.embedding = GPTEmbeddingWithCache(vocab_size, context_size, embedding_size)
self.transformers = [GPTTransformerWithCache(embedding_size, heads) for _ in range(blocks)]
self.output = GPTOutput(embedding_size, vocab_size)
super().__init__([self.embedding] + self.transformers + [self.output])
# 推理函数，传递缓存
def forward(self, x: Tensor, use_cache=False, caches=None, offset=0):
x = self.embedding(x, offset)
# 新缓存列表，每个 Transformer 模块有自己的缓存
new_caches = []
for i, layer in enumerate(self.transformers):
# 当前模块的缓存
cache = caches[i] if caches is not None else None
x, new_cache = layer(x, use_cache, cache)
# 更新缓存
new_caches.append(new_cache)
return self.output(x), new_cachesGPT 缓存训练器¶
GPT 缓存训练器 GPTTrainerWithCache 的 generate 函数：
首次生成：无缓存，计算全部词元；
后续生成：有缓存，只计算新词元；
缓存重置：如果序列长度超过上下文尺寸，重置缓存。
class GPTTrainerWithCache(GPTTrainer):
def generate(self, dataset, prompt, steps=256, use_cache=True):
self.model.eval()
tokens = dataset.encode(prompt)
# 缓存
caches = None
with Tensor.no_grad():
for _ in range(steps):
if use_cache and caches is not None and len(tokens) <= dataset.context_size:
# 有缓存且未超出上下文尺寸，只使用最新的词元
feature = Tensor([tokens[-1:]])
offset = len(tokens) - 1
else:
# 无缓存，使用整个上下文尺寸
feature = Tensor([tokens[-dataset.context_size:]])
offset = 0
caches = None
# 推理，同时更新缓存
logits, caches = self.model(feature, use_cache, caches, offset)
# 采样
last_logits = logits.data[0, -1]
exp = np.exp(last_logits - np.max(last_logits))
probs = exp / np.sum(exp)
token = np.random.choice(len(probs), p=probs)
tokens.append(token)
return dataset.decode(tokens)BATCH_SIZE = 16
CONTEXT_SIZE = 256
EMBEDDING_SIZE = 384
HEADS = 6
BLOCKS = 6DATA_FILE = "../../../../data/tinyshakespeare.txt"
MODEL_FILE = "../../../../data/tinyshakespeare-gpt.npz"
dataset = CharDataset(DATA_FILE, BATCH_SIZE, CONTEXT_SIZE)
model = GPTWithCache(dataset.vocab_size, CONTEXT_SIZE, EMBEDDING_SIZE, HEADS, BLOCKS)
loss_fn = GPTLoss()
optimizer = AdamWOptimizer(model.parameters)
trainer = GPTTrainerWithCache(model, loss_fn, optimizer)
trainer.load(MODEL_FILE)无缓存¶
np.random.seed(42)
start = time.time()
print(trainer.generate(dataset, prompt="ROMEO:", use_cache=False))
print(f"{time.time() - start:.8f} seconds")ROMEO:
When I pray.
CAMILLO:
In my liest, I have done drunk.
CAMILLO:
Ay, since the speak, my doubt, fainners!
I am too, as to a kind in honourity,
Can a very disentent till take from here other
out, strike than 'Which children choistones
That would how have m
4.93294787 seconds
有缓存¶
np.random.seed(42)
start = time.time()
print(trainer.generate(dataset, prompt="ROMEO:", use_cache=True))
print(f"{time.time() - start:.8f} seconds")ROMEO:
When I pray.
CAMILLO:
In my liest, I have done drunk.
CAMILLO:
Ay, since the speak, my doubt, fainners!
I am too, as to a kind in honourity,
Can a very disentent till take from here other
out, strike than 'Which children choistones
That would how have m
0.53253603 seconds
本章，我们实现了键值缓存：把已生成词元的键和值缓存起来。新词元只计算和缓存的注意力。生成每个词元的时间复杂度从 降到 ，生成速度大幅提升。
课后练习¶
调整最大生成长度
steps，观察无缓存和有缓存的速度差异。是否生成长度越长，缓存的优势越明显？