前幾天一路從 GPU、Memory Hierarchy 講到 GEMM 和 Tensor Core。

但真的跑 LLM 時，GPU 並不是只做矩陣乘法。

今天先把一個 Llama-style Transformer Layer 攤開來看：

看起來是一個 Layer，實際上卻是很多不同 operator 串起來的 pipeline。

哪些地方是 GEMM？

Transformer 裡的 Linear Layer，最後通常都會變成矩陣乘法。

例如：

1. Q / K / V Projection
2. Output Projection
3. MLP 的 Gate / Up / Down Projection

這些就是前幾天研究 GEMM 和 Tensor Core 的地方。

但不是所有東西都是 GEMM

其他操作可以先粗略分成：

所以一個 Transformer Layer 的速度，不只取決於 GEMM 算得多快。

資料怎麼搬、Softmax 和 RMSNorm 怎麼做，以及中間結果要不要反覆寫回 GPU memory，也都會影響效能。

Source Code 一行，不一定是一個 Kernel

例如 source code 裡可能只有：

output = attention(q, k, v)

但底下其實可能包含：

QKᵀ
→ Scale / Mask
→ Softmax
→ Attention × V

反過來，推論引擎也可能把多個操作 fusion 成更少的 GPU kernels。

所以讀 LLM Infra code 時，可以同時看三層：

Model Layer
    ↓
Framework Operator
    ↓
GPU Kernel

今天的結論

一個 Transformer Layer 可以先記成：

GEMM 提供主要計算
Elementwise / Reduction 把計算串起來
Data Movement 負責把資料送到下一步

今天先建立這張 operator map。

明天再用 profiler 實際看一次：同一個 Layer 的時間，到底花在哪些 operations 上。

Reference

• Meta Llama reference implementation