模型最终不是以"Transformer""ResNet"的形式在硬件上运行,而是被拆成数百到数千个底层算子。AI Infra 的第一步,就是把模型结构还原成可执行的算子图,并理解每个算子的真实执行代价。
你看到的模型结构(论文中的图):
┌─────────────────────────────────────────┐
│ Transformer Block │
│ │
│ x → RMSNorm → Attention → + → │
│ ↓ ↓ │
│ RMSNorm → SwiGLU → + → out │
└─────────────────────────────────────────┘
实际上在 ncnn param 文件中(每个方框 = 一个算子):
输入 [seq, 1024]
│
├─ RMSNorm(rmsn_196) ────→ Split ──→ Gemm(Q) ──→ Reshape ──→ RMSNorm(QK-Norm)
│ └─→ Gemm(K) ──→ Reshape ──→ RMSNorm(QK-Norm)
│ └─→ Gemm(V) ──→ Reshape
│ ↓
│ Permute → RotaryEmbed(RoPE)
│ ExpandDims → Tile (GQA expand)
│ ↓
│ SDPA (with KV Cache)
│ ↓
│ Permute → Reshape → Gemm(O)
│ ↓
├─── Add (残差 1) ←──────────────────────┘
│
├─ RMSNorm ──→ Split ──→ Gemm(gate) ──→ Swish
│ └─→ Gemm(up)
│ ↓
│ Mul(gate × up) → Gemm(down)
│ ↓
├─── Add (残差 2) ←───────────┘
│
输出
数一下: 1个 Transformer Block = 约 35 个 ncnn 算子
Qwen3-0.6B 28 层 = 28 × 35 ≈ 980 + 辅助 = 1017 个算子!
深度学习算子
│
┌──────────┬──────────┼──────────┬──────────┐
│ │ │ │ │
逐元素运算 归约运算 矩阵运算 空间运算 序列运算
│ │ │ │ │
Add,Mul,Div Sum,Mean MatMul Conv2D Attention
Exp,Log Max,Min GEMM Pooling RoPE
ReLU,Swish Softmax InnerProduct Interp Embedding
Sigmoid L2Norm EinSum Padding Sampling
│
┌─────────┼─────────┐
│ │ │
归一化 激活函数 张量操作
│ │ │
BatchNorm ReLU Reshape
LayerNorm GELU Permute
RMSNorm SwiGLU Concat/Split
GroupNorm Sigmoid Tile/Gather
特点:
- 输入输出同形
- 每个元素独立计算,天然并行
- FLOPs 很低(每个元素 1~10 次运算)
- 但内存访问 = 读整个输入 + 写整个输出
瓶颈来源:
┌──────────────────────────────────────┐
│ 为什么逐元素算子可能比 GEMM 还慢? │
│ │
│ Add [1024, 1024]: │
│ 计算: 1024² = 1M 次加法 ≈ 0.001ms │
│ 内存: 读 4MB + 写 4MB ≈ 0.1ms │
│ → 99% 的时间在等内存! │
│ │
│ MatMul [1024, 4096] × [4096, 1024]: │
│ 计算: 2 × 1024² × 4096 ≈ 8.6G ops │
│ 内存: 读 16MB + 写 4MB ≈ 0.3ms │
│ → 大部分时间在计算 │
└──────────────────────────────────────┘
优化策略:
1. 算子融合: Add → ReLU → Add 3个kernel → 1个kernel
减少中间结果的读写
2. 向量化: SIMD 一次处理 4/8/16 个元素
3. 内存布局: 确保连续访问(coalesced access)
特点:
- 沿指定维度将多个值合并为一个
- Reduce dimension 上的元素有依赖关系
- 需要并行归约算法(tree reduction)
Softmax: 归约算子中最复杂的
Step 1: max(x) — 沿维度归约找最大值(数值稳定)
Step 2: exp(x − max) — 逐元素(无依赖)
Step 3: sum(exp(x − max)) — 归约求和(有依赖)
Step 4: exp(x − max) / sum — 逐元素除(无依赖)
归约步骤通常串行,是瓶颈所在。
这是神经网络的最核心计算——占了 80%+ 的计算量。
MatMul: C[M,N] = A[M,K] × B[K,N]
计算量: 2 × M × N × K FLOPs
访存量: (M×K + K×N + M×N) × sizeof(dtype) bytes
Tiling 的重要性:
大矩阵无法一次性放入 cache
→ 切成小块,每个块完整计算完再换下一个
→ 最大化数据复用,减少 cache miss
GPU 上的优化层次:
Level 1: 全局内存 → 共享内存 (tiling by warp)
Level 2: 共享内存 → 寄存器 (tile by thread)
Level 3: Tensor Core 指令 (每个 cycle = 一个 16×16 块)
为什么 1×1 卷积比 3×3 卷积更适合 GPU?
1×1: 等价于每个空间位置的 MatMul
→ 可以直接用 GEMM kernel
3×3: im2col 膨胀 9 倍内存
→ 或者用直接卷积(访存不规则)
LayerNorm / RMSNorm 在现代 Transformer 中每层至少出现 2 次:
计算步骤:
RMSNorm: mean(x²) → rsqrt → scale
LayerNorm: mean(x) → var(x) → normalize → scale + shift
每个元素的 FLOPs 极少(5~10 FLOPs)
但需要:
- 读整个输入(归约求 mean)
- 写整个输出
→ memory-bound 算子
融合机会:
RMSNorm + Linear: 可以融合!
先归一化再做矩阵乘法 → 在读取数据的同时完成归一化
ncnn 和 vLLM 都有这类融合优化
计算强度 = FLOPs / (字节访问量)
FLOPs: 浮点运算次数
字节访问量: 读输入 + 写输出 + 读权重的总字节数
例子——Qwen3 的 Q 投影 (GEMM):
输入: [1, 1024] (decode 阶段)
权重: [1024, 2048]
输出: [1, 2048]
FLOPs = 2 × 1 × 2048 × 1024 = 4,194,304 ≈ 4.2M
字节访问 = (1024 + 2048 + 1024×2048) × 4 bytes
= (1024 + 2048 + 2,097,152) × 4
≈ 8.4 MB
计算强度 = 4.2M / 8.4M ≈ 0.5 FLOPs/Byte
对比 A100 的 roofline:
峰值算力: 312 TFLOPS (FP16 Tensor Core)
内存带宽: 2039 GB/s
Roofline 拐点: 312T / 2039G ≈ 153 FLOPs/Byte
0.5 << 153 → 严重 memory-bound!
(decode 阶段,batch=1 的 GEMM)
例子——Prefill 阶段的 Q 投影:
输入: [512, 1024] (更大的 batch)
权重: [1024, 2048]
输出: [512, 2048]
FLOPs = 2 × 512 × 2048 × 1024 ≈ 2.1G
字节访问 ≈ (512×1024 + 2048×1024 + 512×2048) × 4 ≈ 14 MB
计算强度 = 2100M / 14M ≈ 150 FLOPs/Byte
150 ≈ 153 拐点 → compute-bound!
# 基于实际 profiling 的分类
OPERATOR_PROFILE = {
# Memory-bound (< 10 FLOPs/Byte)
"memory_bound": {
"Add/Mul/Div": "逐元素 → 每个元素只做 1 次运算,读+写",
"ReLU/GELU/Swish": "逐元素激活 → 同逐元素",
"LayerNorm/RMSNorm": "需要归约统计,但仍是 memory 主导",
"Dropout": "仅 mask 操作",
"Reshape/Permute": "零 FLOPs,纯内存操作",
"Concat/Split": "零 FLOPs,纯内存操作",
"RoPE": "逐元素旋转,memory-bound",
},
# Intermediate (10-100 FLOPs/Byte)
"intermediate": {
"Softmax": "归约+exp,内存和计算相当",
"Conv3×3 (small batch)": "权重复用不够,偏向 memory",
"SwiGLU": "Mul+Swish 都是逐元素,但跟在 GEMM 后可融合",
},
# Compute-bound (> 100 FLOPs/Byte)
"compute_bound": {
"MatMul (大矩阵)": "每个权重元素复用多次",
"Conv1×1 (large batch)": "等价于 MatMul",
"Prefill Attention": "QK^T 是大矩阵乘法",
"Conv3×3 (large batch)": "batch 够大时权重复用率上升",
},
}GPU 上每个 kernel launch 有固定开销:
- CUDA: ~5-10 μs per launch
- Vulkan: ~10-50 μs per launch
- CPU: <1 μs (线程同步开销)
Qwen3-0.6B decoder 有 1017 个算子。
如果 GPU 每个 launch 10 μs:
1017 × 10 μs ≈ 10 ms 纯 launch 开销
对比:Prefill 总共可能只需 50 ms
→ 20% 的时间浪费在 launch 上!
这就是为什么需要算子融合:
融合后算子数: 1017 → ~500 → 5 ms launch 开销
单个小算子(如 Add)的时间分解:
CPU 端: GPU 端:
┌─ 准备参数 0.5 μs ─┐
│─ 设置 descriptor 0.3 μs │
│─ 提交到 queue 0.2 μs │ ┌─ GPU 接收 1 μs ─┐
│─ CPU 继续执行 0.1 μs │ │─ 读输入 5 μs │
│ │ → │─ 计算 0.5 μs │
│ Total: ~10 μs │ │─ 写输出 4 μs │
└──────────────────────────┘ └─ Total: ~10.5 μs ─┘
Kernel 执行本身只需 0.5 μs!
但 launch 开销 + 内存读写 = 20 μs
→ 97.5% 的时间花在"非计算"上!
从 ncnn param 文件分析 Qwen3-0.6B 的 1017 个算子:
算子类型 数量 占比 主要用途
────────────────────────────────────────────
Gemm 168 16.5% Q/K/V/O 投影 + gate/up/down
RMSNorm 86 8.5% 注意力和 MLP 前的归一化 (含 QK-Norm)
BinaryOp (Add) 84 8.3% 残差连接 + 其他二元运算
Permute 84 8.3% 维度重排
Reshape 112 11.0% 头拆分/合并
RotaryEmbed 56 5.5% RoPE 位置编码
SDPA 28 2.8% 缩放点积注意力
Split 84 8.3% 特征分支
ExpandDims 56 5.5% GQA 头扩展
Tile 56 5.5% GQA 头复制
Swish 28 2.8% SwiGLU 激活
Mul (BinaryOp) 28 2.8% SwiGLU gate×up
Input/Output 60+ 5.9% 输入输出/缓存节点
其他 47 4.6% Clip, Concat, ...
────────────────────────────────────────────
总计 ~1017 100%
1. Gemm (168 个, 16.5%)
所有带权重的操作
每层: Q(2048→1024) + K(1024→1024) + V(1024→1024) + O(2048→1024)
+ gate(1024→3072) + up(1024→3072) + down(3072→1024)
= 7 个 per layer × 24 层 (有效 attention 层)
2. Reshape (112 个, 11.0%)
几乎全是零 FLOPs 的维度重排
主要用于头拆分: Gemm 输出是展平的 [heads×dim, seq]
Reshape 变成 [dim, heads, seq]
3. RMSNorm (86 个, 8.5%)
每层 3 个(pre-attention, QK-Norm×2, pre-MLP)
+ 最终输出归一化
4. Split (84 个, 8.3%)
每层至少 2 个 Split:
- 将 RMSNorm 输出分给 Q/K/V 三条路径
- 将残差分支分给主路径和 skip connection
5. Permute (84 个, 8.3%)
NCHW ↔ NHWC 之间的维度重排
为 SDPA 准备正确的输入格式
Qwen3-0.6B 单层 Decoder, decode 阶段, batch=1:
算子 FLOPs 内存访问 估算时间(CPU) 估算时间(GPU)
────────────────────────────────────────────────────────────────────
RMSNorm×3 ~50K ~24 KB ~15 μs ~5 μs
Gemm×7 ~60M ~30 MB ~800 μs ~50 μs
Reshape×4 ~0 ~0 (view) ~1 μs ~1 μs
Permute×3 ~0 ~0 (reorder) ~2 μs ~2 μs
RotaryEmbed×2 ~10K ~8 KB ~10 μs ~5 μs
ExpandDims×2 ~0 ~0 ~1 μs ~1 μs
Tile×2 ~0 ~8 KB ~5 μs ~3 μs
SDPA×1 ~5M ~500 KB ~200 μs ~30 μs
Mul×1 ~3K ~12 KB ~5 μs ~3 μs
Swish×1 ~3K ~12 KB ~5 μs ~3 μs
Add×3 ~3K ~12 KB ~10 μs ~5 μs
────────────────────────────────────────────────────────────────────
Total/层 ~65M ~31 MB ~1.05 ms ~110 μs
28 层总计 ~1.8G ~870 MB ~29 ms ~3.1 ms
关键观察:
- CPU 上每层 ~1ms,28 层 ~29ms → ~34 tokens/s
- GPU 上每层 ~110μs,28 层 ~3.1ms → ~320 tokens/s (理论)
- 实际 GPU 更慢:因为多次 kernel launch 和 HBM 访问开销
同一个数学运算 y = W × x,不同场景选不同实现:
场景 实现 硬件
───────────────────────────────────────────────────────────
CPU, 大矩阵 OpenBLAS/MKL x86 AVX512
CPU, 小矩阵 Direct loop ARM NEON
GPU, batch>1 cuBLAS CUDA
GPU, batch=1 自定义 kernel CUDA (小矩阵优化)
Vulkan compute shader Mobile GPU
端侧 CPU 量化 GEMM INT8 ARM NEON
Apple Silicon Accelerate/MPS Apple GPU
"选哪个实现"这个决策 = Kernel Dispatch
// ncnn 中的算子 dispatch 逻辑(以 Convolution 为例)
// src/layer/convolution.cpp → 根据设备和参数选择实现
int Convolution::create_pipeline(const Option& opt) {
if (opt.use_vulkan_compute) {
return new Convolution_vulkan(...); // Vulkan GPU
}
// CPU 路径
#if NCNN_X86
if (use_winograd && kernel == 3 && stride == 1)
return new Convolution_3x3_winograd(...); // x86 Winograd
if (kernel == 1)
return new Convolution_1x1(...); // 1×1 特殊路径
#endif
#if NCNN_ARM
if (use_winograd && kernel == 3)
return new Convolution_3x3_winograd_arm(...);
#endif
// 通用回退
return new Convolution_im2col_gemm(...); // im2col + GEMM 通用
}从 ncnn param 文件中选出 Layer 0 的算子序列(约 35 个算子),为每个算子标注:
算子名称 | 输入 Shape | 输出 Shape | FLOPs | 内存访问 | Compute/Memory/Layout
对 Q 投影 GEMM:
batch=1: FLOPs / Bytes = ?
batch=32: FLOPs / Bytes = ?
batch=256: FLOPs / Bytes = ?
从 memory-bound 变为 compute-bound 的拐点 batch 是多大?
文件: src/layer/sdpa.cpp
任务:
1. 画出 SDPA 的 6 个输入的数据流(Q, K_cur, V_cur, mask, past_K, past_V)
2. 标出 KV Cache concat 发生的位置
3. 标出 Softmax 计算的位置
4. 说明 GQA 的 head 分组如何影响循环
上一篇: 01_system_architecture.md — 六层架构总览 下一篇: 03_llm_inference_pipeline.md — LLM 推理流水线