博主介绍程序喵大人35 - 资深C/C/Rust/Android/iOS客户端开发10年大厂工作经验嵌入式/人工智能/自动驾驶/音视频/游戏开发入门级选手《C20高级编程》《C23高级编程》等多本书籍著译者更多原创精品文章首发gzh见文末记得订阅专栏以防走丢C基础系列专栏C语言基础系列专栏C大佬养成攻略专栏C训练营个人网站咱们前几章把模型结构、算子还有显存账本都盘得差不多了。现在是时候把目光投向这块深色的控制台了——GPU 硬件本身。作为这套前传的最后一章咱们的画风也从之前的车间和白板正式切入到了底层的硬件视角。你看这块 GPU 板卡它的核心其实就两样东西外围提供巨大存储空间的高带宽显存HBM以及中央密密麻麻负责干活的流式多处理器SM。我们在 PyTorch 里用 Python 调用的那一个个高层算子最终都会通过 PyTorch 的分发机制转到 C 底层并启动对应的 CUDA Kernel。这些 Kernel 才是直接驱使 GPU 芯片上数以千计的流式多处理器SM做并行计算的真正载体。读懂了这个过程你以后去看 Nsight Profiler 的性能报告或者听大佬们聊算子融合Kernel Fusion、Flash Attention就不会觉得是在听天书了。一、CPU vs GPU一个核 vs 一万个核一提到跑大模型为什么大家非得用 GPUCPU 就不行吗其实只要对比一下它们的核心构成答案就不言而喻了。你看左边的 CPU它就像是几个超级顶尖的特种兵比如 8 到 16 个核心。每个人的装备都极其豪华主频跑得飞快能到 5GHz特别擅长处理那种逻辑复杂、拐弯抹角、还得排队执行的任务。再看右边的 GPU拿 A100 来说它里面直接塞了将近 7000 个 CUDA 核心这些核心单拎出来看可能反应还有点慢只有 1-2 GHz甚至只会做最简单的乘法和加法。但它强就强在“万人齐发”——面对像深度学习这种全是独立数据运算的任务这几千人同时开工瞬间就能把 CPU 按在地上摩擦。二、GPU 的内存层级HBM / SRAM / 寄存器既然 GPU 算得这么快那数据喂得跟不上怎么办这就逼着 GPU 搞出了极其严苛的“内存层级金字塔”。你看最底下的 HBM高带宽显存它就是咱们常说的“80GB 显存”。所有的权重、激活值平时都老老实实地待在这里。它容量虽然大但离计算核心相对较远存取起来其实是有些拖沓的。往上一层是每个 SM 专属的 Shared Memory类似 SRAM。它的容量极其可怜只有几十 MB 甚至更少但速度快得惊人。算子优化的核心战役往往就在这层打响。比如大名鼎鼎的 Flash Attention说白了就是想尽办法在这一层把活干完死活不愿意去底层慢吞吞的 HBM 里来回搬数据。到了最顶尖的寄存器Registers那就更小更猛了它是直接和计算单元绑定的每个线程独享速度快到飞起。理解了这个金字塔你就抓住了 GPU 优化的七寸。三、Kernel 是什么一段在 GPU 上并行跑的代码说了半天硬件咱们来看看让硬件动起来的指令——Kernel。你可以把 Kernel 简单理解为一段写成 C 或者 CUDA 的函数。但它的牛逼之处在于它的执行方式你看右边这张网格图。当 CPU 下达一条 Kernel 启动指令时GPU 会瞬间召唤出成百上千个线程Thread。这些线程被高度组织化、模块化地编排。其中最基本的调度单位是 Warp线程束每个 Warp 包含固定的 32 个线程。在硬件层面这 32 个线程共享同一个指令发射器在同一时刻跑同一条指令。这种物理层面的同步执行机制被称为步调一致Lockstep。多个 Warp 组合在一起构成 Block线程块。同一个 Block 内部的线程由于物理距离极近可以使用 Shared Memory 进行极高带宽的组内数据共享与互助而所有 Block 集合起来就组成了最外层的 Grid线程网格。当 Kernel 运行的时候这上万个线程同时执行一模一样的代码指令。每个线程利用自己独一无二的线程 IDThread ID去定位和计算对应的输入数据。这就是传说中的单指令多线程SIMTSingle Instruction Multiple Threads模式。正因为 Warp 中的 32 个线程绑定在一起执行当代码中出现复杂的if-else分支时不同线程走不同的路径会导致分支分化Warp Divergence。硬件在此时会采取串行化的方式来依次执行各分支路径从而使得计算资源出现闲置。编写极致性能的 CUDA Kernel 往往需要保持 Warp 内部执行路径的高度一致。这种暴力的并行分发模式是大模型得以流畅运转的物理基石。四、算子怎么变成 Kernel框架在背后做的事那我们在写 PyTorch 的时候又是怎么触发这些 Kernel 的呢其实框架在背后帮我们把脏活累活全干了。你看这条分层路径最左边是你写的一行极简的 Python 代码比如torch.matmul。一敲回车PyTorch 内部的分发系统就会迅速接管它会偷偷看看你的数据是在 CPU 还是 GPU 上用的是 FP16 还是 BF16。确认完身份后它会去仓库里翻出一段早就写好的、极致优化过的 CUDA Kernel可能是 NVIDIA 官方的 cuBLAS也可能是别的库。最后这段代码才会被真正扔到硬件层面去并行执行。我们平时讨论的底层优化比如用 Triton 写算子其实都是在跳过表层的 Python直接在 Kernel 这一层抠细节。五、矩阵乘法为什么是 GPU 的爱回过头来看为什么深度学习里的算子能如此完美地契合 GPU 这套并行逻辑关键就在于——深度学习的大部分计算全都是矩阵乘法Linear、Attention 基本上全是它。你看中间这个拆解图算一个输出矩阵 C它里面的每一个小格子只和 A 的某一行以及 B 的某一列有关和其他格子完全独立。这意味着什么意味着这几万个输出格子可以直接分配给 GPU 里的几万个线程大家不用互相商量不用等对方算完只管闷头算自己的那一块就行。这种天然零依赖的数据结构简直就是为 GPU 量身定制的“甜点”。六、Memory-bound vs Compute-bound瓶颈在哪既然跑在硬件上总会遇到天花板。在 AI Infra 领域咱们通常把算子分为两大类瓶颈也就是大名鼎鼎的 Roofline 模型。一类叫 Compute-bound算力瓶颈比如巨大无比的矩阵乘法。这种算是“幸福的烦恼”因为这意味着你的显存带宽跟得上GPU 的算力被你确确实实地压榨到了极限。另一类叫 Memory-bound显存带宽瓶颈比如之前讲的 LayerNorm、激活函数。它们的特点是计算量极小但数据搬运量很大。GPU 算得太快结果全卡在等数据从 HBM 传过来的路上了。这种时候计算核心其实是在无奈地摸鱼。遇到这类算子我们的终极优化思路就是——合并它们减少读写 HBM 的次数。七、一次 Forward 在 GPU 上的时间轴把上面所有的理论串起来咱们最后来看看一次真正的 Forward Pass在时间轴上到底长什么样。你看这张像心电图一样的时间线。从左到右一个个 Kernel 排着队在 GPU 上轰炸。你能明显看到Linear 和 Attention 这类 Compute-bound 算子霸占了最宽、最长的时间条而那些 Norm 和 Softmax只占了零星的一点点绿色短条。但别小看这些短条每次启动一个 Kernel 都是有隐形开销的。如果让这些小算子频繁起停时间全浪费在跑腿上了。所以现在的工程大神们天天在做的就是把这些零碎的小 Kernel 揉成一个大 KernelKernel Fusion让这根时间轴变得越紧凑越好。八、总结你现在可以读 Transformer 系列了好了到这里咱们这套【AI专栏】图解深度学习-AI infra工程师必知必会的前传就正式完结了。咱们退后一步看看你现在手里的这套“心智模型”第一步你知道了模型就是权重加结构第二步你看懂了流淌在里面的 Tensor 和 Shape第三步你认识了各个算子工位的成本第四步你理清了训练和推理那两笔截然不同的显存账本最后咱们又一起钻进了 GPU搞懂了 Kernel 是怎么一回事。有了这套严丝合缝的基础认知接下来你再去看那些动辄谈论 KV Cache 管理、多卡分布式并行、甚至是手写算子的硬核文章绝对不会再有一头雾水的感觉了。现在带上这套知识地图你可以自信地跨入正传【AI专栏】图解Transformer从token到推理引擎的大门了码字不易欢迎大家点赞关注评论谢谢
