☰
从零构建大模型:GPU高性能编程与Kernel融合实战
2026/10/22 3:12:55 网站建设 项目流程

最近把一门编号为 CS336 的课程跟读到了第六讲,课程主题很直白——从零构建大模型。前几讲还在聊数据、分词器、模型结构、训练稳定性,第六讲画风突然往下沉,整整一大节课都在讲 GPU 高性能编程和 Kernel 融合。说实话,如果你是第一次接触这类内容,很容易被“Kernel 融合”这个术语劝退,但它其实是在大模型训练和推理中反复出现的性能核心。这一讲解决的是最现实的问题:为什么同样的模型结构,在你的机器上跑得慢吞吞,换个代码写法却能吞吐翻倍。


1. 为什么第六讲要从模型算法扎到硬件层

1.1 算子的“中间落地税”

大多数做大模型的人,日常工作里其实和 CUDA 内核离得很远,更多是在框架层面搭网络、调超参、写数据集 pipeline。但到了训练和推理优化阶段,就会遇到一个绕不开的事实:框架里的一个算子,并不等于 GPU 上的一次执行。每个算子启动时,都有内核启动开销、张量读写开销。更隐蔽的是,算子与算子之间只要产生了中间张量,这些张量往往要被写回显存,再在下一次内核中被读出来。

这就引出了第六讲反复强调的概念:算子的“中间落地税”。举个例子,一个很常见的表达式residual + activation(linear(x)),如果拆成三次内核执行,第一次线性层算出输出,第二次做激活函数,第三次做残差相加。中间多出来的两次写读,全都在全局内存上完成。全局内存带宽是 GPU 上最稀缺的资源之一,很多时候计算单元在等数据,而不是在计算。

1.2 从前几讲到第六讲,知识地图的转折

课程前几讲更像是在“搭积木”,关注模型能不能训练起来、损失能不能正常下降。第六讲则在回答另一个问题:模型训练起来了,能不能高效地训练?这一讲的重点,是教大家建立“硬件执行坐标”——你的代码经过编译器、驱动、调度器之后,到底在 GPU 的哪一层硬件上跑了多少个线程、读写了几次内存、触发了多少次缓存未命中。

我自己的感受是,这一讲实际上是从“程序员的模型视角”切换到“硬件的执行视角”。你不一定需要成为内核开发专家,但需要知道性能瓶颈在哪里、哪些融合手段收益最大、如何用工具定位问题。对于后续要自己做分布式训练、长序列 attention、大规模推理优化的人来说,这是最重要的一块地基。


2. GPU 高性能编程的三个底层认知

2.1 并行不是“更多核”,而是“海量轻量线程”

很多人刚学 GPU 编程时,容易把它想象成“CPU 的高核数版本”。但第六讲一开始就纠正了这个印象。CPU 的单核能力很强,有缓存、分支预测、乱序执行;而 GPU 的设计思路是:我有成千上万的轻量线程,靠大量并发来隐藏内存延迟。

一个很常见的类比是快递分拣仓库。CPU 像一个全能主管,一个问题从头到尾想得很周到;GPU 则像几百个分拣员同时上阵,每个人只处理手边很小的一包快件。分拣员之间虽然同步成本高,但一个人卡住了,其他分拣员也不会闲着太久。GPU 上也是这样,一个线程等待显存读取时,调度器会切换到另一批可运行的线程,把等待时间填满。

2.2 内存访问模式决定性能上限

这讲里最重要的结论之一:内存访问是否“合并”,影响可能远大于算法本身的运算量。所谓合并访问,就是一个线程束(warp)里的 32 个线程,在同一个时间点去读连续的内存地址。GPU 会把这种连续访问合并成少数几次大的内存事务,效率极高。反过来,如果每个线程跳着读,或者线程之间地址分散,一次内存请求会拆成很多次,带宽利用率断崖式下降。

在深度学习算子中,最典型的现象是:按行读取矩阵很快,按列读取矩阵很慢。原因就是行主序存储下,按行读天然连续,按列读会让同一个 warp 里的线程访问跨度很大的地址。第六讲用这个例子解释了一个重要原则:先看数据在显存里的物理布局,再决定线程和索引怎么映射。

2.3 占用率、线程块与寄存器

讲到线程块大小设置,很多人会凭感觉选 128、256、512。第六讲给出了一个更系统的视角——占用率。占用率指的是一个流多处理器上活跃的线程束数量和硬件最大线程束数量的比值。占用率并不是越高越好,但它直接影响 GPU 隐藏延迟的能力。

一个常见误区是:为了减少循环,让每个线程处理大量数据,结果把寄存器用量顶到了硬件上限,导致能够同时驻留的线程块数量变少,最终性能反而不降。实际调优中,比较稳妥的做法是:先设一个 256 或 512 的块大小,再根据寄存器用量和共享内存用量观察占用率,最后用性能分析器确认究竟是不是被全局内存访问拖住了。第六讲反复强调:不要凭直觉优化,先量化瓶颈。


3. Kernel 融合到底“融合”掉了什么

3.1 一次显存读写的旅行

Kernel 融合,简单说就是把多个 GPU 内核函数合并成一个。它看似只是减少了内核启动次数,但真正的收益往往来自减少全局内存读写。我们要分清楚两类开销:一类是启动开销,对于一个几微秒的小内核来说不可忽略;另一类是中间张量读写,这才是大模型场景里更关键的部分。

举个例子,x -> softmax -> dropout -> linear这三个步骤。非融合版本中,softmax 的结果要写回显存,dropout 再读出来并写回,linear 再读。每个中间张量都经过一次 HBM 往返。融合版本则可以在一个内核中,把数据块读进寄存器或共享内存,做完所有操作,最后只把最终结果写回。相当于把一个流程从“每个环节都进一次库房”变成“一条流水线从头干到尾”。

3.2 以 Softmax 为例的融合过程

Softmax 是 Attention 里最常出现、也最适合拿来理解融合的算子。非融合版本里,为了数值稳定性,一般要分三步:先找一行里的最大值,用最大值做指数衰减,再求和并归一化。这就意味着必须把输入矩阵扫描一遍,得到最大值;再扫描一遍,算指数和求和;再扫描一遍,做除法。三次扫描,至少三次全局内存访问,中间还要存一个指数结果。

融合版本可以在一个内核里完成:把一行数据加载到片上,先求最大值,再做减法和指数,再求和,最后归一化。整行数据只需要从全局内存读一次,最终结果写一次。如果是在某些流程中当场读取输入、当场写出,可以减少中间的多次读取。

3.3 内存受限与计算受限:判断融合收益

课程里给了一个判断融合收益的思维模型:你的内核到底是计算受限还是内存受限?用“计算量除以访存量”这个比例来判断。如果计算量很小但访存量很大,说明消耗的时间主要花在搬运数据上,这种内核融合收益最大。反之,如果内核本来就已经把计算单元塞满,融合再怎么减少中间读写,收益都不会特别明显。

我在一个实际算子中深切体会过这个区别。一个逐元素操作x * scale + bias,理论计算量少得可怜,数据规模很大,瓶颈几乎完全在显存带宽。把它们合并成一个内核,启动次数从三次变一次,HBM 读写总量减少一大半,时间自然肉眼可见地下降。但如果是一个矩阵乘法为主的模块,计算强度本来就高,融合其他小算子的收益更多来自减少启动开销,而不是面访存。


4. 把第六讲的关键点落到代码里

4.1 从零写一个融合 Softmax 内核

第六讲配套的代码练习,通常会要求你亲手实现一个融合内核。下面是一个简化但结构完整的融合 Softmax 内核示例,我用的是 Triton 这种面向 GPU 编程的领域专用语言来写,目的是展示融合思路,而不是纠结底层细节。

import triton import triton.language as tl @triton.jit def fused_softmax_kernel(X, Y, stride, n_cols, BLOCK_SIZE: tl.constexpr): row = tl.program_id(0) col_offsets = tl.arange(0, BLOCK_SIZE) mask = col_offsets < n_cols x = tl.load(X + row * stride + col_offsets, mask=mask, other=-float("inf")) row_max = tl.max(x, axis=0) x = x - row_max num = tl.exp(x) denom = tl.sum(num, axis=0) y = num / denom tl.store(Y + row * stride + col_offsets, y, mask=mask)

这段代码核心就做了几件事:按块读取一行,块内归约求最大值,减去最大值再做指数,求和归一化,最后写回。从全局内存读一次、写一次,中间全部在片上完成。这里的BLOCK_SIZE要大于等于列数,并用mask处理边界。如果你直接把这段代码编译成 GPU 内核,会发现它与一个多步骤框架实现相比,省掉了两到三次中间张量写回。

我在写这类内核时,踩过的第一个坑是把other=-float("inf")理解成“用一个极大负值占位”。这个占位值对于 max 归约没问题,但在 exp 计算时,-inf - row_max还会是-inf,exp 之后变成 0,对分母也没有影响。这是经过数值设计后的合理写法,而不是随手选的小数,否则会对结果造成不可控的偏差。

4.2 用编译器自动获得部分融合

不是每个人都需要手写内核。第六讲里也明确提到,现代深度学习框架已经能在图编译层面帮我们做算子融合。例如在 PyTorch 这种框架里,你可以写一个很普通的计算函数,然后交给图编译器处理:

import torch @torch.compile def f(x, w, b): h = torch.matmul(x, w) return torch.nn.functional.gelu(h + b)

在普通模式下,这个函数可能会被拆成矩阵乘、加法、GELU 几次内核调用。经过图编译器处理之后,编译器会分析算子之间的依赖关系,把能够合并的逐元素操作合并到矩阵乘之前或之后执行,减少中间张量的落盘。对于很多训练和推理脚本来说,这是一条“性价比最高的路”:不用改业务逻辑,不用手动写内核,只要确保代码里的张量操作足够简单,图编译器就能帮上忙。

不过也要注意,图编译对这种计算图的融合能力有一定限制。如果代码里大量使用不可编译的原生 Python 控制流、动态形状、或者频繁调用外部扩展,编译器可能保守地回退到逐算子执行。课程里的建议是:把图编译当成第一层融合手段,手写内核当成第二层手段,两者配合。

4.3 基准测试不要只看“变快了”

第六讲花了不少时间讲性能测试方法。最核心的一点是:不要直接看一次运行时间,因为 GPU 内核第一次启动会有初始化开销,而且运行期间可能混入显存动态分配、驱动编译等噪音。正确做法是先预热,再多次执行,取中位数或最小值,并且要在统计上比较可信。

另一个实用指标是有效带宽:把内核搬运的总字节数除以运行时间。比如一个元素级操作,读写 1 亿个浮点数,大约就是 8 亿字节的读加上 8 亿字节的写。如果运行时间是 2 毫秒,有效带宽大约是 800 GB/s。不同实现之间比较有效带宽,比单纯比较时间更接近真相。

我整理过一张对比表,记录同一表达式在不同实现方式下的表现。这里的数字只是示意,但规律很典型:

实现方式内核启动次数HBM 读写总量相对运行时间
拆成乘法/加法/激活三个内核3 次约 6N 字节约 100%
融合成单内核1 次约 2N 字节约 35%

表中 N 代表一次处理的数据规模。融合版本的时间下降并不只是省掉了启动开销,更大一部分来自中间张量不再被写回再读入。如果你的融合结果只看启动次数,就会低估它的价值。


5. 大模型场景里最值得掌握的融合模式

5.1 Attention:分块在线 Softmax

第六讲在大模型相关融合案例中,重点提到了 Attention 机制。一个常见实现是先算Q @ K^T,得到完整的注意力分数矩阵,再对这个矩阵做 Softmax,然后乘 V。这个流程在序列长度很小的时候没问题,一旦序列长度变大,中间那个N x N的注意力矩阵就会非常大,而且读写它非常昂贵。

融合思路是分块处理:不一次性生成完整的注意力矩阵,而是一次只加载一个块,在块内完成部分 Softmax 和部分加权求和。因为 Softmax 的归一化分母依赖整行,所以要采用“在线 Softmax”的技巧,在遍历过程中持续维护一个 row_max 和 row_sum,每处理一个新块就更新一次。整个过程只需要从全局内存读取 Q、K、V,并最终写回 O,中间矩阵不会完整落盘。这和前面一维 Softmax 的融合逻辑是一脉相承的,差别在于它是二维分块的在线版本。

5.2 LayerNorm/RMSNorm:把归约留在片上

大模型中的每一个 Transformer 块几乎都含有 LayerNorm 或 RMSNorm。从计算公式看,归一化涉及均值、方差、缩放、平移。非融合实现会先扫一遍算均值和方差,再扫一遍归一化,再写回结果。一个简单的归一化,可能带来三到五次全局内存访问。融合实现则在算子内部做两遍或在线统计,然后再逐元素完成归一化,整体上只读一次原始输入、写一次输出。

RMSNorm 比 LayerNorm 更少一个均值计算,只需要求均方根,然后做缩放。这种算子在大模型里出现频率极高,哪怕只省掉一次全局读写,在千亿参数规模下也算得上非常可观的收益。尤其当模型采用并行 Transformer 结构时,归一化核融合能直接减少前向和反向中的大量中间张量。

5.3 残差、激活和线性层:融合红利藏在细节里

大模型中到处都是x + dropout(gelu(linear(x)))这类组合。如果你把它们全拆开,每个算子都会写一次中间张量。而实际上,线性层之后的加法和激活是很典型的“尾部融合”:可以在矩阵乘的同一个内核里把 bias 加上,再把激活函数直接套在输出上,然后和残差相加,最后再决定要不要做 dropout。

反向传播里也有类似机会。比如一个权重更新过程,梯度计算可能同时要算损失对权重、对输入、对中间变量的三个梯度。如果把它们拆成三个独立内核,就要重复读取很多相同的中间结果。融合之后,数据复用率明显提高。第六讲里有一个观点我非常认同:在深度学习性能优化里,真正有用的不是你只会背某个算子,而是你看到表达式时,能意识到哪些中间结果是可以“不落地”的。


6. 实操中的典型问题与排查思路

6.1 线程束发散:能躲就躲,躲不了就分流

GPU 中一个线程束内的线程在执行同一个指令时是相对同步的。如果你的代码里出现if (tid % 2 == 0)这种分支,同一个线程束里一半线程走 A 分支,一半走 B 分支,那 GPU 只能分两次执行,等于浪费了一半线程的吞吐。这在复杂业务逻辑里很难完全避免,但可以通过重新组织数据来减少发散。

一个常用办法是把需要做不同处理的数据提前分组。好比查身份证号码时,如果所有人都排在一个队伍里,遇到一个需要人工审核的人,整个队伍都要停下来;提前把人分成普通通道和人工通道,GPU 里的分支就变成两块连续的、规则的数据段。这个思想在后面的数据预处理和 kernel 索引设计中经常用到。

6.2 共享内存和 Bank 冲突

当内核使用共享内存时,最容易出现的一个性能问题是 bank conflict。GPU 的共享内存通常被划分成一系列 bank,不同线程同时访问不同 bank 时效率最高,如果多个线程访问同一个 bank 里的不同地址,请求就会被串行化。举个例子,如果线程 i 访问数组中的s[2 * i],在常见的硬件配置下,两个线程之间会形成冲突。

解决办法很多:可以在共享内存数组里“打 padding”,也就是给每一行多留一个空位,改变线程访问的 bank 分布;也可以重新设计索引,让同一个 warp 内的地址尽量均匀地落在不同 bank 上。这种问题用眼睛很难看出来,最好在性能分析工具里直接查看共享内存冲突计数。第六讲的建议是:当你的内核已经排除了全局内存瓶颈,但共享内存操作占比很高时,下一步就要查 bank conflict。

6.3 profiler 视角:先分内存受限还是计算受限

排查性能问题时,我一般先看两个关键读数:DRAM 吞吐率,和计算单元吞吐率。如果 DRAM 吞吐率非常高,接近硬件的峰值带宽,同时计算单元利用率却不理想,那这是一个明确的内存受限内核,下一步优化方向是减少数据搬运、优化访问模式、做更大力度的融合。反过来,如果计算单元吞吐很高,DRAM 吞吐很低,那就要去检查有没有低效计算指令、不必要的类型转换,或者线程内分支。

我还遇到过一个非常典型的现象:内核的 SM 占用率很高,但 DRAM 吞吐也不低,最终测下来发现大量时间花在线程依赖和等待全局内存返回值上。这种情况下单纯提高占用率已经没有意义,因为瓶颈已经不在活跃线程数量,而在于每个线程等待数据的链路太长。正确的做法是检查每个线程访问数据的局部性,尽量多用寄存器或共享内存减少对全局内存的重复请求。


7. 写在课后的话

这一讲给我的最大触动,不是某个具体内核怎么写,而是“性能”这件事在大模型项目里究竟意味着什么。你可以在 Python 层做很多巧妙的算法改进,但如果底层内核每次都在重复搬运冗余数据,那所有高层优化都会被打折扣。反过来,即使你用最简单的算法,只要能把握好数据复用、合并访问和 Kernel 融合这三件事,也能获得非常可观的性能提升。

如果让我给刚接触这部分内容的人一个建议,那就是:先别急着啃完整的 CUDA 手册,把一个简单的逐元素算子从头到尾拆开,看它在非融合和融合两种实现下的有效带宽差异,再逐渐过渡到归一化、Softmax、Attention 这些真实模型组件。这套思路在第六讲之后,几乎是所有大模型性能工作的通用主线。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询