1. 从一次深夜的算力账单说起为什么我们要死磕GEMM凌晨三点我被一条短信惊醒。不是家人的问候而是云服务商发来的月度账单预警。屏幕上那个刺眼的数字让我瞬间睡意全无。作为一个正在训练一个中等规模语言模型大约70亿参数的团队负责人我清楚地知道这笔费用的核心不是存储不是网络而是计算——具体来说是成千上万张GPU卡上那些名为“矩阵乘加”GEMM的操作在疯狂燃烧。这可能是所有涉足大模型领域的工程师、研究员或创业者迟早要面对的现实大模型的训练和推理本质上是一场由GEMM主导的算力消耗战。你看到的每一次流畅的对话生成每一次精准的文本续写背后都是天文数字般的矩阵乘法在支撑。不理解GEMM就无法理解大模型的成本构成不优化GEMM就无法在效率和成本之间找到平衡点。很多人初入此领域会沉迷于模型架构的创新、损失函数的设计这固然重要。但当你真正把模型“跑”起来尤其是在试图将其产品化、追求响应速度和降低推理成本时你会发现所有的优雅理论最终都要落地为芯片上高效执行的数值计算。而GEMM正是这个转换过程中最核心、最耗时的部分。它就像一个黑洞吞噬了绝大部分的算力资源。因此本文将彻底拆解GEMM与大模型算力消耗之间的深层逻辑。我们不会停留在“GEMM很重要”的结论上而是深入到计算图、硬件执行和内存访问的层面解释为什么是GEMM以及我们可以从哪些切实可行的方向去优化它。无论你是正在学习大模型的学生还是面临线上服务性能压力的工程师理解这些内容都将帮助你更清醒地看待手中的算力资源并做出更明智的技术决策。2. 大模型的计算图GEMM是如何成为绝对主角的要理解算力消耗首先得看清模型在“计算”什么。现代大模型无论是Transformer、GPT还是LLaMA系列其核心计算模块都可以被高度抽象为几种固定模式。当我们打开一个训练框架如PyTorch的Profiler性能分析器对前向传播和反向传播过程进行跟踪时一幅由各种算子构成的计算图便呈现出来。而其中耗时最长的“热点”Hotspot区域几乎无一例外地被几种特定形式的GEMM所占据。2.1 Transformer架构中的GEMM三重奏以最经典的Transformer Decoder Block为例我们可以将其核心计算分解为三个主要的GEMM操作注意力机制中的Q/K/V投影与输出投影Linear Layers操作输入序列的每个token向量维度为d_model会分别通过三个独立的线性层即全连接层投影得到查询Query、键Key、值Value矩阵。这本质上是三个独立的[seq_len, d_model]矩阵与[d_model, d_head * num_heads]权重矩阵的乘法。计算量每个投影都是一次大规模的GEMM。在自注意力中QK^T又是一次[seq_len, d_head]与[d_head, seq_len]的矩阵乘计算复杂度为O(seq_len^2 * d_head)这是Transformer处理长序列时算力暴涨的根源。最后注意力加权求和后的结果再通过一个输出投影层另一个线性层变回d_model维度。前馈网络Feed-Forward Network, FFN操作这是每个Transformer块中参数最密集的部分。通常由两个线性层组成中间夹着一个激活函数如GeLU、Swish。第一个线性层将维度从d_model扩展到d_ff通常是d_model的4倍第二个线性层再投影回d_model。计算量这两个线性层是典型的“宽矩阵”乘法。例如对于一个70亿参数的模型d_model可能是4096d_ff就是16384。那么一次FFN前向传播就包含两次巨大的GEMM[batch*seq, 4096] [4096, 16384]和[batch*seq, 16384] [16384, 4096]。其计算强度计算量/内存访问量通常比注意力部分更高。嵌入层与输出层Embedding LM Head操作将输入的token ID转换为向量嵌入层以及将最终的隐藏层向量转换为词汇表上的概率分布输出层通常与嵌入层权重共享。计算量这通常是一个查找表Look-up操作加上一个GEMM。虽然查找表本身不是GEMM但输出层的计算[batch*seq, d_model] [d_model, vocab_size]是一个极其“瘦高”的矩阵乘法vocab_size可达数万甚至数十万对内存带宽极其敏感优化不当极易成为瓶颈。注意这里有一个关键点容易被忽略。在框架中一个nn.Linear层的前向传播在底层会被分解为一次矩阵乘法GEMM和一次向量加法Bias Add。而GEMM的优化收益远大于Bias Add因此我们的焦点始终在GEMM上。2.2 计算强度与“内存墙”问题为什么GEMM如此受关注除了它出现频率高更因为它具有极高的计算强度潜力。计算强度定义为总计算操作数FLOPs / 总内存访问字节数Bytes。一个理想的、完全在高速缓存Cache中进行的GEMM可以做到每从内存中读取一个数据元素例如一个FP16的数就进行大量的乘加运算FMA。现代GPU的峰值算力TFLOPs非常高但内存带宽TB/s的提升相对缓慢。这就形成了“内存墙”。计算密集型Compute-Bound当计算强度足够高使得计算单元持续忙碌等待数据的时间很少性能受限于GPU的峰值算力。优化良好的大型、规整GEMM如FFN中的矩阵乘通常处于这种状态。内存密集型Memory-Bound当计算强度低计算单元大部分时间在等待数据从慢速的全局内存中加载进来性能受限于内存带宽。小矩阵乘法、不规则访问、以及像嵌入层输出那样的大维度投影很容易陷入这种状态。大模型中的GEMM优化很大程度上就是在与“内存墙”搏斗想方设法将内存密集型操作转化为计算密集型操作或者至少减少不必要的内存访问。3. GEMM算力消耗的定量分析你的FLOPs花在了哪里理解了GEMM在哪里出现我们还需要定量地知道它到底有多“贵”。这对于估算训练成本、推理延迟和硬件选型至关重要。3.1 如何估算一个模型的前向传播FLOPs对于一个线性层Y X W b忽略偏置b其浮点运算次数FLOPs可以精确计算。矩阵X形状为[M, K]权重W形状为[K, N]输出Y为[M, N]。一次乘加运算FMA通常被视为2次浮点运算一次乘法一次加法。计算Y中的一个元素y_ij需要K次乘加运算即2*KFLOPs。整个矩阵共有M * N个元素因此总FLOPs为2 * M * N * K。以一个具体的70亿参数模型的一层FFN为例假设d_model4096,d_ff16384,batch_size1,seq_len1024。第一个线性层X: [1*1024, 4096],W: [4096, 16384]。FLOPs 2 * 1048576 * 4096 * 16384 ≈1.76e14(176 TFLOPs)。第二个线性层X: [1048576, 16384],W: [16384, 4096]。FLOPs同样约为1.76e14(176 TFLOPs)。 仅仅这一层FFN的一次前向传播就需要约352 TFLOPs的算力。而一次完整的训练迭代包含前向、反向传播计算量大约是前向的2-3倍和优化器更新其算力消耗可想而知。3.2 训练与推理的算力消耗差异训练和推理对GEMM的消耗模式有显著不同训练阶段算力消耗巨兽需要完整的反向传播和优化器步骤。反向传播中为了计算权重梯度需要做大量的GEMM运算例如计算dY/dW本身就是另一个X^T dY的矩阵乘法。内存消耗更高需要存储中间激活值用于反向传播这对显存构成巨大压力。激活重计算Activation Checkpointing等技术就是用额外的计算重新执行前向来换取显存节省这又增加了GEMM的执行次数。关注吞吐量通常使用大批次Large Batch来尽可能压满GPU的算力追求的是每秒处理的样本数Tokens/s。推理阶段自回归解码的延迟挑战在生成文本时模型以自回归方式运行每次只预测下一个token。这意味着大部分GEMM操作的Mbatch*seq维度中seq是逐步增长的。早期步骤的矩阵非常“瘦”M小属于内存密集型难以利用GPU的并行能力导致延迟高、硬件利用率低。关注延迟与成本追求单个请求的响应速度Time to First Token, TTFT和每美元所能处理的token数Tokens per Dollar。计算图静态化推理框架如TensorRT, ONNX Runtime会提前将动态图优化、融合成静态的Kernel其中GEMM的优化是重中之重。3.3 硬件利用率的现实峰值算力的“水分”芯片厂商宣传的峰值算力如A100的312 TFLOPS for FP16是在最理想、最规整的GEMM下测得的。现实中由于以下原因实际利用率Achieved FLOPS往往大打折扣矩阵形状不理想小批量Small Batch、短序列Short Sequence导致矩阵维度M,N过小无法隐藏内存延迟也无法启动足够多的线程块填充SM流多处理器。数据布局不佳例如权重矩阵不是内存连续访问友好的布局如默认的Row-Major导致访存效率低下。精度转换开销混合精度训练中在FP16和FP32之间的类型转换、缩放Scaling操作会引入额外开销。Kernel启动开销频繁启动大量小规模的GEMM Kernel其启动开销本身可能成为瓶颈。因此优化GEMM的首要目标就是通过各种手段让实际的GEMM运算尽可能接近芯片的理论峰值。4. 核心优化策略一算法与模型架构层面的“降本增效”在动手写代码或调库之前最高效的优化往往来自于算法和模型设计本身。这相当于在源头减少“需求”。4.1 稀疏化与剪枝让矩阵“瘦身”如果权重矩阵W中有大量接近零的值那么这些值参与的乘加运算对结果贡献微乎其微却消耗着等量的算力。稀疏化旨在识别并消除这些冗余。结构化剪枝整行、整列或整块地移除权重。例如将[4096, 16384]的矩阵剪枝为[4096, 12288]。优点是剪枝后的矩阵仍然是稠密的可以直接使用高度优化的标准GEMM库如cuBLAS。难点在于如何最小化精度损失。非结构化剪枝随机地移除单个权重元素可能达到很高的稀疏率如90%。但产生的矩阵是高度不规则的稀疏矩阵无法使用标准GEMM。需要专门的稀疏矩阵乘法库如cuSPARSE或硬件支持如NVIDIA的稀疏张量核心。只有当稀疏度非常高时其带来的计算节省才能抵消稀疏计算本身的开销。实践心得对于大模型推理2:4结构化稀疏每4个元素中必有2个为零是一个实用的折衷方案。NVIDIA Ampere架构及之后的GPU其张量核心可以直接加速这种模式的稀疏GEMM理论上能获得2倍的吞吐提升。许多推理框架已经支持加载这种格式的模型。4.2 低秩近似用“小矩阵”替代“大矩阵”其核心思想是一个大的稠密矩阵W [M, N]可以用两个小矩阵的乘积来近似W ≈ A [M, r] B [r, N]其中r秩远小于M和N。LoRALow-Rank Adaptation在微调阶段冻结原始大模型权重W只训练两个低秩矩阵A和B。这样前向传播变为Y X W X (A B)。虽然看起来多了一项GEMM但AB的计算量2*M*r*N远小于直接微调W的计算量2*M*N*K这里K是输入维度。LoRA极大地降低了微调的可训练参数量和显存需求。内在思考LoRA之所以有效一个假设是模型在任务适配时权重变化具有“低秩”特性。这启示我们在模型设计时是否可以直接用低秩分解的结构来构建某些层以在推理时获得天然的速度优势一些轻量化模型架构正在探索这条道路。4.3 量化用“轻量级”数据做“重型”计算量化将高精度浮点数如FP32转换为低精度格式如INT8, INT4甚至二进制1-bit。这从三个方面优化GEMM减少内存占用和带宽压力INT8数据大小是FP32的1/4传输同样大小的矩阵带宽需求降至1/4。提升计算吞吐许多硬件如GPU的INT8张量核心执行低精度运算的峰值算力远高于高精度。降低能耗。权重量化Weight-only Quantization仅对权重进行量化激活值保持高精度。在推理时将INT8权重反量化为FP16再计算。这主要节省了模型加载的显存和内存带宽对GEMM计算本身的加速有限因为计算仍在FP16下进行。权激活值量化Weight-Activation Quantization对权重和激活值同时量化。GEMM核心计算在INT8下进行。这能最大程度发挥硬件INT8算力但挑战在于如何校准激活值的动态范围以及处理异常值Outliers精度损失风险更高。GPTQ、AWQ等后训练量化技术通过在少量校准数据上微调寻找最优的量化参数在几乎不损失精度的情况下将大模型量化到INT4甚至更低。这是当前大模型推理部署的标配技术。实操陷阱量化不是简单的类型转换。你需要仔细选择量化方案对称/非对称、校准方法最大最小值、百分位、熵校准以及粒度每张量、每通道、每组。错误的量化会导致模型精度崩溃。务必使用成熟的量化工具包如TensorRT, GPTQ实现 Hugging Face的optimum库并进行严格的评估。5. 核心优化策略二系统与运行时层面的“精打细算”当模型架构确定后我们需要在软件栈层面确保每一个GEMM都能以最高效的方式在硬件上执行。5.1 算子融合减少Kernel启动与内存读写在原始的计算图中一个线性层可能被分解为GEMM、Bias Add、Activation等多个独立的Kernel。每个Kernel都需要启动开销并且需要将中间结果写回全局内存供下一个Kernel读取这产生了巨大的冗余内存流量。融合GEMM Bias Add Activation将这三个操作合并成一个自定义的CUDA Kernel。在这个融合Kernel中GEMM计算出的一个结果元素立即加上偏置再通过激活函数然后才写回全局内存。这消除了中间结果的读写极大提升了数据复用降低了延迟。现代推理引擎如TensorRT, TVM的核心能力之一就是自动进行此类算子融合。FlashAttention极致的融合典范它不是一个简单的算子融合而是对标准Attention计算过程的彻底重写。通过“平铺”Tiling技术将大的QK^T和Softmax计算分解为小块在SRAM共享内存中进行计算并重新排列计算顺序以避免将庞大的中间矩阵[seq_len, seq_len]写回HBM高带宽内存。这完美解决了长序列下的内存瓶颈问题是算法与系统优化结合的巅峰之作。5.2 内存布局与访问优化数据在内存中如何排列直接决定了访问效率。行主序 vs 列主序cuBLAS等库默认期望矩阵是列主序Column-Major。而Python/NumPy/PyTorch默认是行主序Row-Major。虽然库内部会处理转换但隐式的转置操作意味着一次额外的内存遍历和拷贝。在性能关键路径上确保数据布局符合库的期望是基本要求。内存对齐GPU内存访问通常有对齐要求如128字节对齐。确保矩阵数据起始地址和步长stride是对齐的可以使内存访问合并Coalesced一次性读取连续的数据块极大提升带宽利用率。连续内存避免使用非连续的张量视图如transpose()后不调用contiguous()。非连续内存访问会导致缓存效率低下严重时性能下降数倍。在调用GEMM前检查并确保输入矩阵是内存连续的。5.3 选择正确的计算库与后端不要试图自己手写GEMM Kernel除非你是这方面的专家。站在巨人的肩膀上是最明智的选择。cuBLAS / cuBLASLt (NVIDIA)基础且强大。cublasGemmEx和cublasLtMatmul支持混合精度、多种数据类型和算法选择。cublasLt更轻量适合在运行时动态选择最优的算法。CUTLASS (NVIDIA)一个用于编写高性能GEMM和其他线性代数运算的模板库。它提供了模块化的组件允许专家级开发者针对特定的矩阵形状、数据类型和硬件进行极致优化。许多前沿的融合算子如FlashAttention都是基于CUTLASS实现的。oneDNN (Intel)/OpenBLAS在CPU上进行高性能GEMM的首选。框架集成PyTorch通过torch.compile以及背后的TorchInductor可以自动将多个操作融合并生成优化的Kernel。TensorRT则通过解析ONNX模型进行层间融合、精度校准和Kernel自动调优生成高度优化的推理引擎。一个实际调优案例在部署一个BERT模型时我们发现一个线性层的GEMM性能不佳。使用Nsight Systems分析发现该Kernel的“Achieved Occupancy”SM占用率很低。通过检查发现输入矩阵的Batch维度很小M32但K和N很大。我们尝试了两种方案1使用cublasLtMatmul并让其自动寻找最优算法2将多个独立的小GEMM批量打包成一个更大的GEMMtorch.bmm或手动拼接。最终方案2通过增加M维度显著提高了硬件利用率和吞吐量。6. 面向推理的专项优化应对自回归解码的挑战推理尤其是文本生成的自回归解码对GEMM优化提出了独特要求。6.1 批处理与持续批处理静态批处理将多个请求在序列长度维度进行填充Padding后拼接成一个批次。优点是简单能提高GPU利用率。缺点是效率受制于最长的序列短序列存在大量无效计算Padding部分。持续批处理也称为迭代级调度或流式批处理。它动态地管理一个批处理队列当一个请求生成结束遇到EOS token后立即从队列中移出并可能加入新的请求。这极大地提高了GPU利用率和吞吐量。vLLM、TGIText Generation Inference等流行推理服务器都实现了此功能。其核心挑战在于高效地管理变长序列的KV Cache和注意力计算。6.2 KV Cache避免重复计算的“记忆”在自回归解码中第t步的K和V矩阵实际上包含了前t-1步的所有历史键值对。如果不做缓存每一步都需要重新计算所有历史token的K和V计算量是O(t^2)。KV Cache技术将每一步计算出的K_t,V_t存储下来供后续步骤使用。这样每一步的注意力计算只需要计算当前新token的Q_t, K_t, V_t然后与缓存的K_{1:t-1}, V_{1:t-1}进行拼接和计算。这避免了O(t^2)的重复GEMM将计算复杂度降回O(t)。内存挑战KV Cache会消耗大量显存。对于一个L层、H个头、d_head维度、batch_size为B、seq_len为S的模型KV Cache的总大小约为2 * L * B * S * H * d_head * bytes_per_param。优化KV Cache的内存布局和压缩是推理优化的关键。PagedAttentionvLLM提出的革命性技术。它将连续的KV Cache空间划分为固定大小的“块”Block并像操作系统管理内存页一样管理这些块。不同序列的KV Cache可以非连续地存储从而高效处理变长序列并几乎消除因碎片化导致的内存浪费。6.3 推测解码用“猜测”换取并行自回归解码的根本瓶颈在于其串行性。推测解码试图打破这一点。核心思想用一个小的、快速的“草稿模型”一次性生成多个候选token一个推测序列。然后用原始大模型“验证模型”并行地对整个候选序列进行验证。如果验证通过则一次性接受多个token大幅提升解码速度。如何并行验证这正是GEMM的用武之地。验证模型需要并行地计算整个候选序列的注意力。这要求模型能够处理一个批次中每个序列具有不同“前缀候选”的复杂注意力掩码。高效的实现需要精心设计Kernel以处理这种不规则的计算模式。局限性如果草稿模型的预测准确率不高会导致大量候选被拒绝验证的计算就浪费了。因此草稿模型的质量和与主模型的对齐程度至关重要。7. 工具链与实战如何定位并优化你的GEMM瓶颈理论说了这么多最终要落地。以下是一个可操作的性能分析与优化流程。7.1 性能剖析找到真正的热点PyTorch Profiler最易上手的工具。使用torch.profiler记录训练或推理过程重点关注GPU Kernel视图。你会看到volta_fp16_s884gemm_fp16_128x128_ldg8_f2f这样的Kernel名称它们就是cuBLAS的GEMM Kernel。查看它们的耗时占比、GPU利用率GPU Utilization和Tensor Core利用率Tensor Core Utilization。Nsight Systems更底层的系统级性能分析工具。它可以提供时间线上每个Kernel的精确耗时、内存拷贝、CUDA API调用等信息。特别有用的是它可以告诉你一个Kernel是“Compute Bound”还是“Memory Bound”这是优化方向的根本指引。Nsight Compute针对单个Kernel的微观架构分析工具。如果你怀疑某个特定的GEMM Kernel性能不佳可以用它进行深入分析查看SM占用率、内存带宽利用率、指令发射效率等数百个指标找到具体的瓶颈点。7.2 一个完整的优化迭代案例假设我们正在优化一个线上对话模型的推理服务发现平均响应延迟过高。Profiling使用PyTorch Profiler对单个请求进行跟踪。发现耗时最高的操作是lm_head输出投影层的GEMM且该Kernel被频繁调用每次生成一个token调用一次。分析lm_head的权重矩阵形状是[d_model, vocab_size]。在自回归解码中M batch_size * 1每次只处理最新token的隐藏状态N vocab_size很大如50000K d_model。这是一个典型的“瘦高”矩阵乘法M很小极度内存密集型GPU利用率极低。优化方案方案A量化对lm_head的权重进行INT8量化。由于它是内存密集型减少权重内存带宽能直接提升性能。使用GPTQ进行精度无损的INT4量化效果更佳。方案B算子融合检查计算图发现lm_head之后紧跟着log_softmax。使用TensorRT构建引擎将GEMM Bias Add LogSoftmax融合成一个Kernel减少中间数据读写。方案C批处理优化服务端采用持续批处理将多个用户请求的动态批次中同一解码步的lm_head计算进行批处理。即使每个请求的M1多个请求的M累加起来也能形成一个更“胖”的矩阵提高计算强度。方案D算法替代对于超大词表可以考虑使用“自适应Softmax”或“采样Softmax”来避免计算整个词表的概率但这会改变模型行为需要重新训练或微调。实施与验证我们选择实施方案AINT4量化和方案B算子融合。使用TensorRT构建量化后的融合引擎。重新Profiling发现lm_head阶段的耗时下降了约70%整体延迟显著改善。7.3 经验与陷阱不要过早优化首先确保你的代码和模型是正确的。在优化之前建立一个准确的性能基线。量化不是银弹量化会引入精度损失。必须使用验证集评估量化后模型的准确率如困惑度PPL、任务准确率。对于敏感任务可能需要部分量化或更精细的量化策略。融合的边界并非所有算子都能无缝融合。复杂的控制流如条件判断、动态形状等会增加融合的难度。需要依赖成熟的编译器如TorchDynamo, TensorRT来处理。硬件特性不同架构的GPU如Ampere, Hopper其张量核心、内存体系结构不同最优的GEMM参数如Thread Block大小也可能不同。利用库的自动调优功能如cublasLt的启发式搜索往往比手动调参更有效。优化GEMM是一场贯穿大模型生命周期的持久战。从模型架构设计时的稀疏化、低秩思想到训练时的混合精度与梯度累积策略再到推理时的量化、融合与解码优化每一环都深刻影响着最终的算力消耗和用户体验。理解其背后的原理熟练运用 profiling 工具定位瓶颈并灵活组合各种优化技术是我们从“算力账单”的焦虑中走向从容掌控的必经之路。 大模型算力黑洞：拆解GEMM优化，从原理到实战降本增效
1. 从一次深夜的算力账单说起为什么我们要死磕GEMM凌晨三点我被一条短信惊醒。不是家人的问候而是云服务商发来的月度账单预警。屏幕上那个刺眼的数字让我瞬间睡意全无。作为一个正在训练一个中等规模语言模型大约70亿参数的团队负责人我清楚地知道这笔费用的核心不是存储不是网络而是计算——具体来说是成千上万张GPU卡上那些名为“矩阵乘加”GEMM的操作在疯狂燃烧。这可能是所有涉足大模型领域的工程师、研究员或创业者迟早要面对的现实大模型的训练和推理本质上是一场由GEMM主导的算力消耗战。你看到的每一次流畅的对话生成每一次精准的文本续写背后都是天文数字般的矩阵乘法在支撑。不理解GEMM就无法理解大模型的成本构成不优化GEMM就无法在效率和成本之间找到平衡点。很多人初入此领域会沉迷于模型架构的创新、损失函数的设计这固然重要。但当你真正把模型“跑”起来尤其是在试图将其产品化、追求响应速度和降低推理成本时你会发现所有的优雅理论最终都要落地为芯片上高效执行的数值计算。而GEMM正是这个转换过程中最核心、最耗时的部分。它就像一个黑洞吞噬了绝大部分的算力资源。因此本文将彻底拆解GEMM与大模型算力消耗之间的深层逻辑。我们不会停留在“GEMM很重要”的结论上而是深入到计算图、硬件执行和内存访问的层面解释为什么是GEMM以及我们可以从哪些切实可行的方向去优化它。无论你是正在学习大模型的学生还是面临线上服务性能压力的工程师理解这些内容都将帮助你更清醒地看待手中的算力资源并做出更明智的技术决策。2. 大模型的计算图GEMM是如何成为绝对主角的要理解算力消耗首先得看清模型在“计算”什么。现代大模型无论是Transformer、GPT还是LLaMA系列其核心计算模块都可以被高度抽象为几种固定模式。当我们打开一个训练框架如PyTorch的Profiler性能分析器对前向传播和反向传播过程进行跟踪时一幅由各种算子构成的计算图便呈现出来。而其中耗时最长的“热点”Hotspot区域几乎无一例外地被几种特定形式的GEMM所占据。2.1 Transformer架构中的GEMM三重奏以最经典的Transformer Decoder Block为例我们可以将其核心计算分解为三个主要的GEMM操作注意力机制中的Q/K/V投影与输出投影Linear Layers操作输入序列的每个token向量维度为d_model会分别通过三个独立的线性层即全连接层投影得到查询Query、键Key、值Value矩阵。这本质上是三个独立的[seq_len, d_model]矩阵与[d_model, d_head * num_heads]权重矩阵的乘法。计算量每个投影都是一次大规模的GEMM。在自注意力中QK^T又是一次[seq_len, d_head]与[d_head, seq_len]的矩阵乘计算复杂度为O(seq_len^2 * d_head)这是Transformer处理长序列时算力暴涨的根源。最后注意力加权求和后的结果再通过一个输出投影层另一个线性层变回d_model维度。前馈网络Feed-Forward Network, FFN操作这是每个Transformer块中参数最密集的部分。通常由两个线性层组成中间夹着一个激活函数如GeLU、Swish。第一个线性层将维度从d_model扩展到d_ff通常是d_model的4倍第二个线性层再投影回d_model。计算量这两个线性层是典型的“宽矩阵”乘法。例如对于一个70亿参数的模型d_model可能是4096d_ff就是16384。那么一次FFN前向传播就包含两次巨大的GEMM[batch*seq, 4096] [4096, 16384]和[batch*seq, 16384] [16384, 4096]。其计算强度计算量/内存访问量通常比注意力部分更高。嵌入层与输出层Embedding LM Head操作将输入的token ID转换为向量嵌入层以及将最终的隐藏层向量转换为词汇表上的概率分布输出层通常与嵌入层权重共享。计算量这通常是一个查找表Look-up操作加上一个GEMM。虽然查找表本身不是GEMM但输出层的计算[batch*seq, d_model] [d_model, vocab_size]是一个极其“瘦高”的矩阵乘法vocab_size可达数万甚至数十万对内存带宽极其敏感优化不当极易成为瓶颈。注意这里有一个关键点容易被忽略。在框架中一个nn.Linear层的前向传播在底层会被分解为一次矩阵乘法GEMM和一次向量加法Bias Add。而GEMM的优化收益远大于Bias Add因此我们的焦点始终在GEMM上。2.2 计算强度与“内存墙”问题为什么GEMM如此受关注除了它出现频率高更因为它具有极高的计算强度潜力。计算强度定义为总计算操作数FLOPs / 总内存访问字节数Bytes。一个理想的、完全在高速缓存Cache中进行的GEMM可以做到每从内存中读取一个数据元素例如一个FP16的数就进行大量的乘加运算FMA。现代GPU的峰值算力TFLOPs非常高但内存带宽TB/s的提升相对缓慢。这就形成了“内存墙”。计算密集型Compute-Bound当计算强度足够高使得计算单元持续忙碌等待数据的时间很少性能受限于GPU的峰值算力。优化良好的大型、规整GEMM如FFN中的矩阵乘通常处于这种状态。内存密集型Memory-Bound当计算强度低计算单元大部分时间在等待数据从慢速的全局内存中加载进来性能受限于内存带宽。小矩阵乘法、不规则访问、以及像嵌入层输出那样的大维度投影很容易陷入这种状态。大模型中的GEMM优化很大程度上就是在与“内存墙”搏斗想方设法将内存密集型操作转化为计算密集型操作或者至少减少不必要的内存访问。3. GEMM算力消耗的定量分析你的FLOPs花在了哪里理解了GEMM在哪里出现我们还需要定量地知道它到底有多“贵”。这对于估算训练成本、推理延迟和硬件选型至关重要。3.1 如何估算一个模型的前向传播FLOPs对于一个线性层Y X W b忽略偏置b其浮点运算次数FLOPs可以精确计算。矩阵X形状为[M, K]权重W形状为[K, N]输出Y为[M, N]。一次乘加运算FMA通常被视为2次浮点运算一次乘法一次加法。计算Y中的一个元素y_ij需要K次乘加运算即2*KFLOPs。整个矩阵共有M * N个元素因此总FLOPs为2 * M * N * K。以一个具体的70亿参数模型的一层FFN为例假设d_model4096,d_ff16384,batch_size1,seq_len1024。第一个线性层X: [1*1024, 4096],W: [4096, 16384]。FLOPs 2 * 1048576 * 4096 * 16384 ≈1.76e14(176 TFLOPs)。第二个线性层X: [1048576, 16384],W: [16384, 4096]。FLOPs同样约为1.76e14(176 TFLOPs)。 仅仅这一层FFN的一次前向传播就需要约352 TFLOPs的算力。而一次完整的训练迭代包含前向、反向传播计算量大约是前向的2-3倍和优化器更新其算力消耗可想而知。3.2 训练与推理的算力消耗差异训练和推理对GEMM的消耗模式有显著不同训练阶段算力消耗巨兽需要完整的反向传播和优化器步骤。反向传播中为了计算权重梯度需要做大量的GEMM运算例如计算dY/dW本身就是另一个X^T dY的矩阵乘法。内存消耗更高需要存储中间激活值用于反向传播这对显存构成巨大压力。激活重计算Activation Checkpointing等技术就是用额外的计算重新执行前向来换取显存节省这又增加了GEMM的执行次数。关注吞吐量通常使用大批次Large Batch来尽可能压满GPU的算力追求的是每秒处理的样本数Tokens/s。推理阶段自回归解码的延迟挑战在生成文本时模型以自回归方式运行每次只预测下一个token。这意味着大部分GEMM操作的Mbatch*seq维度中seq是逐步增长的。早期步骤的矩阵非常“瘦”M小属于内存密集型难以利用GPU的并行能力导致延迟高、硬件利用率低。关注延迟与成本追求单个请求的响应速度Time to First Token, TTFT和每美元所能处理的token数Tokens per Dollar。计算图静态化推理框架如TensorRT, ONNX Runtime会提前将动态图优化、融合成静态的Kernel其中GEMM的优化是重中之重。3.3 硬件利用率的现实峰值算力的“水分”芯片厂商宣传的峰值算力如A100的312 TFLOPS for FP16是在最理想、最规整的GEMM下测得的。现实中由于以下原因实际利用率Achieved FLOPS往往大打折扣矩阵形状不理想小批量Small Batch、短序列Short Sequence导致矩阵维度M,N过小无法隐藏内存延迟也无法启动足够多的线程块填充SM流多处理器。数据布局不佳例如权重矩阵不是内存连续访问友好的布局如默认的Row-Major导致访存效率低下。精度转换开销混合精度训练中在FP16和FP32之间的类型转换、缩放Scaling操作会引入额外开销。Kernel启动开销频繁启动大量小规模的GEMM Kernel其启动开销本身可能成为瓶颈。因此优化GEMM的首要目标就是通过各种手段让实际的GEMM运算尽可能接近芯片的理论峰值。4. 核心优化策略一算法与模型架构层面的“降本增效”在动手写代码或调库之前最高效的优化往往来自于算法和模型设计本身。这相当于在源头减少“需求”。4.1 稀疏化与剪枝让矩阵“瘦身”如果权重矩阵W中有大量接近零的值那么这些值参与的乘加运算对结果贡献微乎其微却消耗着等量的算力。稀疏化旨在识别并消除这些冗余。结构化剪枝整行、整列或整块地移除权重。例如将[4096, 16384]的矩阵剪枝为[4096, 12288]。优点是剪枝后的矩阵仍然是稠密的可以直接使用高度优化的标准GEMM库如cuBLAS。难点在于如何最小化精度损失。非结构化剪枝随机地移除单个权重元素可能达到很高的稀疏率如90%。但产生的矩阵是高度不规则的稀疏矩阵无法使用标准GEMM。需要专门的稀疏矩阵乘法库如cuSPARSE或硬件支持如NVIDIA的稀疏张量核心。只有当稀疏度非常高时其带来的计算节省才能抵消稀疏计算本身的开销。实践心得对于大模型推理2:4结构化稀疏每4个元素中必有2个为零是一个实用的折衷方案。NVIDIA Ampere架构及之后的GPU其张量核心可以直接加速这种模式的稀疏GEMM理论上能获得2倍的吞吐提升。许多推理框架已经支持加载这种格式的模型。4.2 低秩近似用“小矩阵”替代“大矩阵”其核心思想是一个大的稠密矩阵W [M, N]可以用两个小矩阵的乘积来近似W ≈ A [M, r] B [r, N]其中r秩远小于M和N。LoRALow-Rank Adaptation在微调阶段冻结原始大模型权重W只训练两个低秩矩阵A和B。这样前向传播变为Y X W X (A B)。虽然看起来多了一项GEMM但AB的计算量2*M*r*N远小于直接微调W的计算量2*M*N*K这里K是输入维度。LoRA极大地降低了微调的可训练参数量和显存需求。内在思考LoRA之所以有效一个假设是模型在任务适配时权重变化具有“低秩”特性。这启示我们在模型设计时是否可以直接用低秩分解的结构来构建某些层以在推理时获得天然的速度优势一些轻量化模型架构正在探索这条道路。4.3 量化用“轻量级”数据做“重型”计算量化将高精度浮点数如FP32转换为低精度格式如INT8, INT4甚至二进制1-bit。这从三个方面优化GEMM减少内存占用和带宽压力INT8数据大小是FP32的1/4传输同样大小的矩阵带宽需求降至1/4。提升计算吞吐许多硬件如GPU的INT8张量核心执行低精度运算的峰值算力远高于高精度。降低能耗。权重量化Weight-only Quantization仅对权重进行量化激活值保持高精度。在推理时将INT8权重反量化为FP16再计算。这主要节省了模型加载的显存和内存带宽对GEMM计算本身的加速有限因为计算仍在FP16下进行。权激活值量化Weight-Activation Quantization对权重和激活值同时量化。GEMM核心计算在INT8下进行。这能最大程度发挥硬件INT8算力但挑战在于如何校准激活值的动态范围以及处理异常值Outliers精度损失风险更高。GPTQ、AWQ等后训练量化技术通过在少量校准数据上微调寻找最优的量化参数在几乎不损失精度的情况下将大模型量化到INT4甚至更低。这是当前大模型推理部署的标配技术。实操陷阱量化不是简单的类型转换。你需要仔细选择量化方案对称/非对称、校准方法最大最小值、百分位、熵校准以及粒度每张量、每通道、每组。错误的量化会导致模型精度崩溃。务必使用成熟的量化工具包如TensorRT, GPTQ实现 Hugging Face的optimum库并进行严格的评估。5. 核心优化策略二系统与运行时层面的“精打细算”当模型架构确定后我们需要在软件栈层面确保每一个GEMM都能以最高效的方式在硬件上执行。5.1 算子融合减少Kernel启动与内存读写在原始的计算图中一个线性层可能被分解为GEMM、Bias Add、Activation等多个独立的Kernel。每个Kernel都需要启动开销并且需要将中间结果写回全局内存供下一个Kernel读取这产生了巨大的冗余内存流量。融合GEMM Bias Add Activation将这三个操作合并成一个自定义的CUDA Kernel。在这个融合Kernel中GEMM计算出的一个结果元素立即加上偏置再通过激活函数然后才写回全局内存。这消除了中间结果的读写极大提升了数据复用降低了延迟。现代推理引擎如TensorRT, TVM的核心能力之一就是自动进行此类算子融合。FlashAttention极致的融合典范它不是一个简单的算子融合而是对标准Attention计算过程的彻底重写。通过“平铺”Tiling技术将大的QK^T和Softmax计算分解为小块在SRAM共享内存中进行计算并重新排列计算顺序以避免将庞大的中间矩阵[seq_len, seq_len]写回HBM高带宽内存。这完美解决了长序列下的内存瓶颈问题是算法与系统优化结合的巅峰之作。5.2 内存布局与访问优化数据在内存中如何排列直接决定了访问效率。行主序 vs 列主序cuBLAS等库默认期望矩阵是列主序Column-Major。而Python/NumPy/PyTorch默认是行主序Row-Major。虽然库内部会处理转换但隐式的转置操作意味着一次额外的内存遍历和拷贝。在性能关键路径上确保数据布局符合库的期望是基本要求。内存对齐GPU内存访问通常有对齐要求如128字节对齐。确保矩阵数据起始地址和步长stride是对齐的可以使内存访问合并Coalesced一次性读取连续的数据块极大提升带宽利用率。连续内存避免使用非连续的张量视图如transpose()后不调用contiguous()。非连续内存访问会导致缓存效率低下严重时性能下降数倍。在调用GEMM前检查并确保输入矩阵是内存连续的。5.3 选择正确的计算库与后端不要试图自己手写GEMM Kernel除非你是这方面的专家。站在巨人的肩膀上是最明智的选择。cuBLAS / cuBLASLt (NVIDIA)基础且强大。cublasGemmEx和cublasLtMatmul支持混合精度、多种数据类型和算法选择。cublasLt更轻量适合在运行时动态选择最优的算法。CUTLASS (NVIDIA)一个用于编写高性能GEMM和其他线性代数运算的模板库。它提供了模块化的组件允许专家级开发者针对特定的矩阵形状、数据类型和硬件进行极致优化。许多前沿的融合算子如FlashAttention都是基于CUTLASS实现的。oneDNN (Intel)/OpenBLAS在CPU上进行高性能GEMM的首选。框架集成PyTorch通过torch.compile以及背后的TorchInductor可以自动将多个操作融合并生成优化的Kernel。TensorRT则通过解析ONNX模型进行层间融合、精度校准和Kernel自动调优生成高度优化的推理引擎。一个实际调优案例在部署一个BERT模型时我们发现一个线性层的GEMM性能不佳。使用Nsight Systems分析发现该Kernel的“Achieved Occupancy”SM占用率很低。通过检查发现输入矩阵的Batch维度很小M32但K和N很大。我们尝试了两种方案1使用cublasLtMatmul并让其自动寻找最优算法2将多个独立的小GEMM批量打包成一个更大的GEMMtorch.bmm或手动拼接。最终方案2通过增加M维度显著提高了硬件利用率和吞吐量。6. 面向推理的专项优化应对自回归解码的挑战推理尤其是文本生成的自回归解码对GEMM优化提出了独特要求。6.1 批处理与持续批处理静态批处理将多个请求在序列长度维度进行填充Padding后拼接成一个批次。优点是简单能提高GPU利用率。缺点是效率受制于最长的序列短序列存在大量无效计算Padding部分。持续批处理也称为迭代级调度或流式批处理。它动态地管理一个批处理队列当一个请求生成结束遇到EOS token后立即从队列中移出并可能加入新的请求。这极大地提高了GPU利用率和吞吐量。vLLM、TGIText Generation Inference等流行推理服务器都实现了此功能。其核心挑战在于高效地管理变长序列的KV Cache和注意力计算。6.2 KV Cache避免重复计算的“记忆”在自回归解码中第t步的K和V矩阵实际上包含了前t-1步的所有历史键值对。如果不做缓存每一步都需要重新计算所有历史token的K和V计算量是O(t^2)。KV Cache技术将每一步计算出的K_t,V_t存储下来供后续步骤使用。这样每一步的注意力计算只需要计算当前新token的Q_t, K_t, V_t然后与缓存的K_{1:t-1}, V_{1:t-1}进行拼接和计算。这避免了O(t^2)的重复GEMM将计算复杂度降回O(t)。内存挑战KV Cache会消耗大量显存。对于一个L层、H个头、d_head维度、batch_size为B、seq_len为S的模型KV Cache的总大小约为2 * L * B * S * H * d_head * bytes_per_param。优化KV Cache的内存布局和压缩是推理优化的关键。PagedAttentionvLLM提出的革命性技术。它将连续的KV Cache空间划分为固定大小的“块”Block并像操作系统管理内存页一样管理这些块。不同序列的KV Cache可以非连续地存储从而高效处理变长序列并几乎消除因碎片化导致的内存浪费。6.3 推测解码用“猜测”换取并行自回归解码的根本瓶颈在于其串行性。推测解码试图打破这一点。核心思想用一个小的、快速的“草稿模型”一次性生成多个候选token一个推测序列。然后用原始大模型“验证模型”并行地对整个候选序列进行验证。如果验证通过则一次性接受多个token大幅提升解码速度。如何并行验证这正是GEMM的用武之地。验证模型需要并行地计算整个候选序列的注意力。这要求模型能够处理一个批次中每个序列具有不同“前缀候选”的复杂注意力掩码。高效的实现需要精心设计Kernel以处理这种不规则的计算模式。局限性如果草稿模型的预测准确率不高会导致大量候选被拒绝验证的计算就浪费了。因此草稿模型的质量和与主模型的对齐程度至关重要。7. 工具链与实战如何定位并优化你的GEMM瓶颈理论说了这么多最终要落地。以下是一个可操作的性能分析与优化流程。7.1 性能剖析找到真正的热点PyTorch Profiler最易上手的工具。使用torch.profiler记录训练或推理过程重点关注GPU Kernel视图。你会看到volta_fp16_s884gemm_fp16_128x128_ldg8_f2f这样的Kernel名称它们就是cuBLAS的GEMM Kernel。查看它们的耗时占比、GPU利用率GPU Utilization和Tensor Core利用率Tensor Core Utilization。Nsight Systems更底层的系统级性能分析工具。它可以提供时间线上每个Kernel的精确耗时、内存拷贝、CUDA API调用等信息。特别有用的是它可以告诉你一个Kernel是“Compute Bound”还是“Memory Bound”这是优化方向的根本指引。Nsight Compute针对单个Kernel的微观架构分析工具。如果你怀疑某个特定的GEMM Kernel性能不佳可以用它进行深入分析查看SM占用率、内存带宽利用率、指令发射效率等数百个指标找到具体的瓶颈点。7.2 一个完整的优化迭代案例假设我们正在优化一个线上对话模型的推理服务发现平均响应延迟过高。Profiling使用PyTorch Profiler对单个请求进行跟踪。发现耗时最高的操作是lm_head输出投影层的GEMM且该Kernel被频繁调用每次生成一个token调用一次。分析lm_head的权重矩阵形状是[d_model, vocab_size]。在自回归解码中M batch_size * 1每次只处理最新token的隐藏状态N vocab_size很大如50000K d_model。这是一个典型的“瘦高”矩阵乘法M很小极度内存密集型GPU利用率极低。优化方案方案A量化对lm_head的权重进行INT8量化。由于它是内存密集型减少权重内存带宽能直接提升性能。使用GPTQ进行精度无损的INT4量化效果更佳。方案B算子融合检查计算图发现lm_head之后紧跟着log_softmax。使用TensorRT构建引擎将GEMM Bias Add LogSoftmax融合成一个Kernel减少中间数据读写。方案C批处理优化服务端采用持续批处理将多个用户请求的动态批次中同一解码步的lm_head计算进行批处理。即使每个请求的M1多个请求的M累加起来也能形成一个更“胖”的矩阵提高计算强度。方案D算法替代对于超大词表可以考虑使用“自适应Softmax”或“采样Softmax”来避免计算整个词表的概率但这会改变模型行为需要重新训练或微调。实施与验证我们选择实施方案AINT4量化和方案B算子融合。使用TensorRT构建量化后的融合引擎。重新Profiling发现lm_head阶段的耗时下降了约70%整体延迟显著改善。7.3 经验与陷阱不要过早优化首先确保你的代码和模型是正确的。在优化之前建立一个准确的性能基线。量化不是银弹量化会引入精度损失。必须使用验证集评估量化后模型的准确率如困惑度PPL、任务准确率。对于敏感任务可能需要部分量化或更精细的量化策略。融合的边界并非所有算子都能无缝融合。复杂的控制流如条件判断、动态形状等会增加融合的难度。需要依赖成熟的编译器如TorchDynamo, TensorRT来处理。硬件特性不同架构的GPU如Ampere, Hopper其张量核心、内存体系结构不同最优的GEMM参数如Thread Block大小也可能不同。利用库的自动调优功能如cublasLt的启发式搜索往往比手动调参更有效。优化GEMM是一场贯穿大模型生命周期的持久战。从模型架构设计时的稀疏化、低秩思想到训练时的混合精度与梯度累积策略再到推理时的量化、融合与解码优化每一环都深刻影响着最终的算力消耗和用户体验。理解其背后的原理熟练运用 profiling 工具定位瓶颈并灵活组合各种优化技术是我们从“算力账单”的焦虑中走向从容掌控的必经之路。