Reference
FLOPs
定义
每秒浮点运算次数, Floating Point Operations per Second
浮点数的加法、乘法、除法。FLOPs 表示一个模型在一次前向计算中需要执行多少次浮点运算
MFLOPs是million FLOPs, GFLOPs是giga FLOPs
矩阵FLOPs

添加图片注释,不超过 140 字(可选)

添加图片注释,不超过 140 字(可选)
线性层FLOPs

添加图片注释,不超过 140 字(可选)
粗略估算

添加图片注释,不超过 140 字(可选)
代码
估算FLOPs的一段代码
fastai course2: 14_augment.ipynb
def _flops(x, h, w): # x是tensor,height, width
if x.dim()<3: return x.numel() # numel返回的是元素数量,比如shape(3,4)返回的是12
if x.dim()==4: return x.numel()*h*w
@fc.patch
def summary(self:Learner):
res = '|Module|Input|Output|Num params|MFLOPS|\n|--|--|--|--|--|\n'
totp,totf = 0,0
def _f(hook, mod, inp, outp):
nonlocal res,totp,totf
nparms = sum(o.numel() for o in mod.parameters())
totp += nparms
*_,h,w = outp.shape
flops = sum(_flops(o, h, w) for o in mod.parameters())/1e6
totf += flops
res += f'|{type(mod).__name__}|{tuple(inp[0].shape)}|{tuple(outp.shape)}|{nparms}|{flops:.1f}|\n'
with Hooks(self.model, _f) as hooks: self.fit(1, lr=1, cbs=SingleBatchCB())
print(f"Tot params: {totp}; MFLOPS: {totf:.1f}")
if fc.IN_NOTEBOOK:
from IPython.display import Markdown
return Markdown(res)
else: print(res)

添加图片注释,不超过 140 字(可选)
FLOPs占比

添加图片注释,不超过 140 字(可选)
在小模型里,multi head attention (MHA)的FLOPs 和 MLP(FFN) 不相上下。
但是在大模型里,大部分都是MLP,这个时候如果一味的optimize MHA是不work的。
Scaling规模扩大
可以指参数,数据集,或者计算量FLOPs

添加图片注释,不超过 140 字(可选)
较难任务只有在FLOPs突破一定规模的时候才能开始有效

添加图片注释,不超过 140 字(可选)
Modular arithmetic,模运算,17 mod5 =2 , 23mod7=2,最常见的就是时钟mod12, 通常对大模型比较难
IPA transliterate,把英文单词转换成音标,需要打模型对发音的理解。
word unscramble, 把英文单词里的字母打乱,让模型重新组合成单词
Persian QA, 波斯语言的问答,波斯语的语库很小,如果能正确问答,说明模型具有泛化能力。
Truthful QA,真实问答,问生鱼片能治病吗?错误是生鱼片是个好东西,正确是不能。
Grounded mapping, 概念映射,问一种生活在非洲的食草动物,脖子很长,答长颈鹿。
Multitask NLU, 多任务语言理解,打开了窗户意味着房间通风了。
word in context,词在上下文中的意思。he banked the plane, 这里bank不是银行而是指转弯。
一些概念
语言模型领域里有一些直觉的东西,很难解释,比如SwiGLU,都是通过实验验证得来的。
直觉的东西很难transfer。
技术性的东西可以transfer,比如transformer如何工作的,比如model parallel里GPU如何被分派,还有mindset上,scaling law很重要之类的。
efficiency
另外,模型的efficiency很重要,2012到2019年间,imagenet efficiency提升了44倍。
而且efficiency对超大模型很重要,因为训练一次的成本很高。
模型的公开程度
closed , 比如4o
open weights,比如deepseek, 不公开模型code,但是给了weights,huggingface上的一些模型,然后有paper作为支持。
open source,比如OLMo, 公开模型code,也给weights,方便重复。
Others
Common Crawl dump 是一个包含了数十亿网页原始内容的公开网络抓取数据集,按月发布,用于训练大型语言模型,包括 GPT、BERT、PaLM
五个模块
可以从五个方面最大化模型efficiency

添加图片注释,不超过 140 字(可选)
tokenizer,是可以把文字转换成integer整数的东西。比如apple转换成1,banana转换成2.
Byte pair encoding
BPE是一种分词方法,它将常见的字符对合并成新的“子词单位”,以减少词汇碎片。
下面是一个例子

添加图片注释,不超过 140 字(可选)
课程上的例子,这个merge的过程可以循环几次。
def train_bpe(string: str, num_merges: int) -> BPETokenizerParams: # @inspect string, @inspect num_merges
Start with the list of bytes of string.
indices = list(map(int, string.encode("utf-8"))) # @inspect indices
merges: dict[tuple[int, int], int] = {} # index1, index2 => merged index
vocab: dict[int, bytes] = {x: bytes([x]) for x in range(256)} # index -> bytes
for i in range(num_merges):
Count the number of occurrences of each pair of tokens
counts = defaultdict(int)
for index1, index2 in zip(indices, indices[1:]): # For each adjacent pair
counts[(index1, index2)] += 1 # @inspect counts
Find the most common pair.
pair = max(counts, key=counts.get) # @inspect pair
index1, index2 = pair
Merge that pair.
new_index = 256 + i # @inspect new_index
merges[pair] = new_index # @inspect merges
vocab[new_index] = vocab[index1] + vocab[index2] # @inspect vocab
indices = merge(indices, pair, new_index) # @inspect indices
return BPETokenizerParams(vocab=vocab, merges=merges)
Transformer architecture

添加图片注释,不超过 140 字(可选)
左边的是一个整体的GPT架构图,右边是灰色transformer block的具体架构图,也就是transformer具体的样子。

添加图片注释,不超过 140 字(可选)
Single-Head
这是一段self-attention的代码(single head),input是image, 来自于fastai 2: 27_attention.ipynb
class SelfAttention(nn.Module):
def __init__(self, ni):
super().__init__()
self.scale = math.sqrt(ni)
self.norm = nn.BatchNorm2d(ni)
self.qkv = nn.Linear(ni, ni*3)
self.proj = nn.Linear(ni, ni)
def forward(self, inp):
n,c,h,w = inp.shape # n:batch size, c:channel, h:height, w: width
x = self.norm(inp).view(n, c, -1).transpose(1, 2) # view merge h,w to hw, transpose makes hw to the center
q,k,v = torch.chunk(self.qkv(x), 3, dim=-1)
s = (q@k.transpose(1,2))/self.scale
x = s.softmax(dim=-1)@v
x = self.proj(x).transpose(1,2).reshape(n,c,h,w)
return x+inp
这里图片是input,shape是n,c,h,w -- n:batch size, c:channel, h:height, w: width
view(n,c-1)把图片shape变成n,c, hw, transpose(1,2)变成了n,hw(seq_len), c,因为transformer永远是在中间这个seq_len上work。
SelfAttention(ni=128)
ni等于c的大小
Multi-Head
多头注意力机制(Multi-Head Attention)是 Transformer 的核心创新之一。它的核心思想是: 并行学习多个注意力模式,比如一个头关注局部信息、一个头关注全局句法、一个头关注实体。
这是一个multi-head version
class SelfAttentionMultiHead(nn.Module):
def __init__(self, ni, nheads):
super().__init__()
self.nheads = nheads
self.scale = math.sqrt(ni/nheads)
self.norm = nn.BatchNorm2d(ni)
self.qkv = nn.Linear(ni, ni*3)
self.proj = nn.Linear(ni, ni)
def forward(self, inp):
n,c,h,w = inp.shape
x = self.norm(inp).view(n, c, -1).transpose(1, 2)
x = self.qkv(x)
x = rearrange(x, 'n s (h d) -> (n h) s d', h=self.nheads)
q,k,v = torch.chunk(x, 3, dim=-1)
s = (q@k.transpose(1,2))/self.scale
x = s.softmax(dim=-1)@v
x = rearrange(x, '(n h) s d -> n s (h d)', h=self.nheads)
x = self.proj(x).transpose(1,2).reshape(n,c,h,w)
return x+inp
这里实现multi-head的关键在于这行
x = rearrange(x, 'n s (h d) -> (n h) s d', h=self.nheads)
把channel给拆成了(h d),每个维度是 d = c // nheads, (channel 在CNN里已经扩了很多,这里只是图片的例子。)
拆后 shape 是 (n * h, s, d)
和单头做一次 softmax 乘积 不同,多头的话可以做nheads个头并行 softmax,最后 concat 回去。表达能力高,可捕捉多个关注模式。
Quadratic blowup(平方级膨胀)
Quadratic blowup 是指 self-attention 的计算/内存复杂度为 O(n^2),当序列长度 n增大时,成本平方级增长。
每一步的计算复杂度:

添加图片注释,不超过 140 字(可选)
关键瓶颈是 Q 与 K 转置相乘, 生成一个 n x n 矩阵
因此,增长量会随着序列大小Quadratic(平方级别n^2)增长, (而不是expotential指数增长哦,2^n,远远快于平方)

添加图片注释,不超过 140 字(可选)
为了防止平方级增长,会有sliding window attention, 还有linear attention
Lower dimension
可以把前面linear的后面的ni*3改成更小的dimension,这里比如加一个小的kqv_dim,代替掉之前的ni x3
class SelfAttentionLowDim(nn.Module):
def __init__(self, c, qkv_dim=32):
super().__init__()
self.scale = math.sqrt(qkv_dim)
self.norm = nn.BatchNorm2d(c)
self.qkv = nn.Linear(c, qkv_dim * 3)
self.proj = nn.Linear(qkv_dim, c) # 把低维输出映射回通道数
好处是可以降低计算复杂度与内存需求。
Lower dimension- group query attention
所有 heads 分为若干组,每组共享 Key/Value 权重。
每个 Query head 拥有独立的 weights Q, 但多个 Query head 共享同一组weights K & V
假设总共有 **16 个 attention heads,**我们设置 **每 4 个 Query head 共享一组 Key/Value。**这就是 GQA ratio = 4:1 的一种配置。
下面是一个chatgpt生成的coding例子
class GroupQueryAttention(nn.Module):
def __init__(self, d_model, n_heads, kv_groups):
super().__init__()
self.n_heads = n_heads
self.kv_groups = kv_groups
self.d_head = d_model // n_heads
self.q_proj = nn.Linear(d_model, d_model)
self.k_proj = nn.Linear(d_model, d_model // kv_groups)
self.v_proj = nn.Linear(d_model, d_model // kv_groups)
self.out_proj = nn.Linear(d_model, d_model)
def forward(self, x):
B, T, _ = x.shape
q = self.q_proj(x).view(B, T, self.n_heads, self.d_head)
k = self.k_proj(x).view(B, T, self.kv_groups, self.d_head * (self.n_heads // self.kv_groups))
v = self.v_proj(x).view(B, T, self.kv_groups, self.d_head * (self.n_heads // self.kv_groups))
# Broadcast k,v across heads within each group...
Low dimension - multi head latent attention
用一个固定数量的latent vectors(例如 256 个)来替代 n 个 queries

添加图片注释,不超过 140 字(可选)
Query 是 m个 learnable latent vectors
输出是 m个注意力结果,可以进一步处理或 decode 回原始序列
一段chatgpt生成的coding例子:
import torch
import torch.nn as nn
class MultiHeadLatentAttention(nn.Module):
def __init__(self, input_dim, latent_dim, num_latents, num_heads):
super().__init__()
self.latents = nn.Parameter(torch.randn(num_latents, latent_dim)) # (m, d)
self.cross_attn = nn.MultiheadAttention(embed_dim=latent_dim, num_heads=num_heads, batch_first=True)
# 如果输入维度与 latent_dim 不一致,先投影
self.input_proj = nn.Linear(input_dim, latent_dim)
def forward(self, x):
"""
x: (batch_size, seq_len, input_dim)
"""
b, n, _ = x.shape
x = self.input_proj(x) # shape: (b, n, latent_dim)
# Expand latent queries to batch: (b, m, d)
latents = self.latents.unsqueeze(0).expand(b, -1, -1)
# Cross-attention: Query = latent, Key/Value = input
out, _ = self.cross_attn(query=latents, key=x, value=x)
return out # shape: (b, num_latents, latent_dim)
# ✨ 测试一下
x = torch.randn(4, 100, 128) # batch=4, seq_len=100, input_dim=128
model = MultiHeadLatentAttention(input_dim=128, latent_dim=64, num_latents=16, num_heads=4)
output = model(x)
print(output.shape) # ➜ (4, 16, 64)
GPU
结构
了解GPU的结构可以帮助写kernel代码
A100的样子,左边是单个 SM(Streaming Multiprocessors),右边是一个128 个 SM 的芯片结构。

添加图片注释,不超过 140 字(可选)
一些关于每个模块的解释

添加图片注释,不超过 140 字(可选)
Warp
Warp 是 GPU 的最小执行单元(32 个线程),由 SM 内的 Warp Scheduler 在每个周期挑选 ready warp 来调度执行,充分隐藏 memory latency 并提升并行效率。
每个时钟周期:
Warp Scheduler 会选一个 ready warp 来执行下一条指令
所谓 ready warp:所有线程不等待 memory、不阻塞、不发散
如果 warp A 在等 memory(如 global memory load),scheduler 会暂时不调它
会切换到另一个 ready warp B
这让 GPU 能在没有 cache 命中时仍保持高吞吐
如果遇到类似这样的case
if x[i] > 0:
y[i] = x[i]
else:
y[i] = -x[i]
GPU 会顺序执行两个分支路径 → 降低并行度(性能下降)
所以 Triton / CUDA 优化中经常使用
mask = cond
val = tl.where(mask, a, b)
Triton
下面是triton的常用命令和对应的GPU结构

添加图片注释,不超过 140 字(可选)
KV Cache
KV cache 是缓存 transformer 中已经计算过的 key 和 value,用来避免重复计算,极大加速自回归生成。
在生成任务(如 GPT)中,模型是 一个 token 一个 token 地生成,例如:
Input: "The cat"
→ 模型生成 "sat"
Input: "The cat sat"
→ 模型生成 "on"
Input: "The cat sat on"
→ 模型生成 "the
每次都重新计算整个序列的 attention(包括前面的所有 token),重复浪费大量算力。

添加图片注释,不超过 140 字(可选)

添加图片注释,不超过 140 字(可选)
每步只计算 1 个 token 的 Q/K/V
K/V 是增量累积、永久缓存,只追加不回算
GPT example
# 每层 transformer 都会存 KV cache
kv_cache = {
"layer_0": {
"k": tensor of shape (batch, num_heads, seq_len, head_dim),
"v": tensor of shape (batch, num_heads, seq_len, head_dim)
},
...
}
在每层中,当前新 token 的 Q 会去 attend 所有历史 K(来自 cache), 得到上下文感知的表示
时间复杂度
- 没有 KV cache:每一步生成都要重新处理整段序列 → O(n²) 时间复杂度
- 有 KV cache:每一步只需处理新 token → O(n) 时间复杂度(更精准地说是 O(1) per step)
Inference
分为两个阶段,左边是prefill phase 预填阶段,右边是decoding phase解码阶段

添加图片注释,不超过 140 字(可选)
预填阶段会先处理prompt,并行处理所有的token,计算每个token的KQV,然后保存到KV Cache
解码阶段会一个一个的生成,生成新的token,用上一个token作为新的输入。

添加图片注释,不超过 140 字(可选)
解码过程中有一些trick:
- 比如用更小的模型(pruning、quantization、distillation)加速生成。
- Speculative decoding,先用一个草稿模型(draft model)一次性预测多个 token, 再用大模型验证这几个 token 是否合理(可以并行打分)