ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

用Triton手写21个Kernel,Qwen3.5推理提速至223 tokens/s

用Triton手写21个Kernel,Qwen3.5推理提速至223 tokens/s 我花了两周时间用 Triton 手写了 21 个 kernel把 Qwen3.5-0.8B 的完整推理链路从 PyTorch 的自动调度里一层层剥出来最终在单张消费级显卡上把生成速度压到了 223 tokens/s。整个过程远没有标题看起来那么光鲜中途遇到过 NaN 乱飞、CUDA Graph 捕获失败、split-K 归约结果对不齐等各种问题。这篇文章是系列第一篇重点讲三件事这 21 个算子到底怎么拆、split-K 在 GEMV 和 RMS Norm 里的两种正确打开方式、以及 CUDA Graph 捕获期的坑怎么绕。适合那些已经用熟 PyTorch、想把手伸到算子层的人也适合刚看完 Triton 官方 tutorial、想拿一个真实模型练手的朋友。读完你能得到一份可以直接抄作业的 kernel 映射表以及一套在 GPU 上调试自定义算子时真正管用的排查思路。1. 为什么是 Qwen3.5-0.8B选型动机与“手写全算子”的可行性判断先说说我为什么选这个模型。0.8B 这个量级很微妙它足够大能说明问题——注意力、FFN、RMS Norm、RoPE、KV Cache 这些主流推理组件全都有不会被“模型太小所以省略细节”糊弄过去它又足够小小到单卡能跑、单算子能肉眼检查复现成本低到可以反复推翻重来。如果一上来就手写 7B 模型的算子光显存和调试时间就能把人拖垮而且性能瓶颈往往在显存带宽而非算子本身写起来反而不容易看清问题。选型另一个判断点是Qwen3.5-0.8B 的架构里没有 MoE、没有专家路由这类让 kernel 复杂度陡增的组件注意力部分用的是 GQAGrouped Query Attention。GQA 意味着 KV heads 数量远小于 Q heads这会让 q_proj / k_proj / v_proj 的矩阵形状严重不对称对后面写 GEMV kernel 是一个非常好的真实案例。模型配置大致是 hidden_size896、24 层 Transformer block、8 个 Q heads、4 个 KV heads、intermediate_size4864词汇表约 15 万。这些数字本身不重要重要的是它们决定了 kernel 里的 BLOCK 大小、split-K 切分方式和显存占用估算。再聊一下“手写全部算子”这个决策的可行性。我给自己定的目标不是替代 cuBLAS也不是跟 vLLM 拼性能而是把model(input_ids)这条链路里的每一步都换成自己能读懂的 Triton kernel然后保证浮点误差在可接受范围内。这个目标听起来大实际拆开就是一张表格的事一个标准 decoder-only Transformer 的推理阶段去掉算子内部的复杂性真正需要实现的独立计算模式其实只有 GEMV、小矩阵乘、elementwise 归约、softmax、gather 这几类。剩下的工作全部是“换数据形状”。所以 21 个算子听起来多但核心代码模式不超过 6 种。我在动手前做了一个很关键的决定把 kernel 按“生成阶段”和“提示阶段”分别写。生成阶段是逐 token 解码所有矩阵乘都退化成 M1 的 GEMV提示阶段是并行 prefilling矩阵乘的 M 等于输入长度。这两个阶段对并行度的要求完全相反混在一起写会两头不讨好。系列第一版我先全部打通生成阶段因为 GEMV 是 split-K 的最佳试验场也是 CUDA Graph 收益最大的地方。2. 21 个算子的全量拆解清单先画地图再动工手写 kernel 之前最重要的一件事不是写代码而是画图。把模型结构翻译成一张算子清单每个算子标清楚输入输出形状、归约维度、是否适合 split-K、是否在 KV Cache 路径上。这张表就是整个项目的施工图后面每写一个 kernel 都是在跟这张表对齐。2.1 算子映射总表从 Transformer 层到 Triton kernel下面是我最终定稿的 21 个算子清单。注意这里统计的是“去重后的独立 kernel”也就是说 24 层 Transformer 共用同一个rms_normkernel循环调用不等于重新实现。编号算子名所属模块核心计算模式是否用 split-K1embedding_lookupEmbeddinggather 取词向量否2rope_embedAttention旋转位置编码 elementwise否3q_proj_gemvAttentionGEMVQ heads 维度切分是4k_proj_gemvAttentionGEMVKV heads 维度切分是5v_proj_gemvAttentionGEMV是6qk_dotAttention1×head_dim 与 S×head_dim 点积否7softmax_maskAttentionmasked softmax含 causal mask否8pv_dotAttention1×head_dim 与 S×head_dim 加权求和否9attn_out_projAttentionGEMV把多头结果合并回 hidden是10residual_addBlock 公共elementwise 加残差否11rms_normBlock 公共按 hidden 维求均方根再 scale是12gate_projMLPGEMV 到 intermediate_size是13up_projMLPGEMV 到 intermediate_size是14silu_mulMLPSiLU 激活后与另一分支相乘否15down_projMLPGEMV 从 intermediate_size 回到 hidden是16final_rms_norm输出层同上归一化到输出是17lm_head_gemv输出层GEMV 到 vocab_size是18splitk_reduce归约工具对 split-K 部分和做合并否19kv_cache_updateKV Cache写入/更新 cache 张量否20cache_maskKV Cache生成长度掩码与位置索引否21token_sampler输出层argmax / top-k 采样否这张表有几点需要解释。首先q_proj / k_proj / v_proj 我拆成了三个独立 kernel而不是像 PyTorch 里那样合并成一个大 GEMV原因是为了应对 GQA 带来的形状不对称q_proj 的输出维度是 num_q_heads × head_dimk/v_proj 是 num_kv_heads × head_dim合并写意味着 kernel 内部要做一堆 mask 和维度判断把简单问题复杂化。拆开之后每个 kernel 的 BLOCK_N 完全不同反而更容易调优。代价是多了一次 kernel launch 的开销但这个开销在 CUDA Graph 捕获后会被压到很小。其次residual_add和rms_norm没有融合。很多高性能实现会把 RMS Norm 前面的残差加法直接并进去减少一次显存读写。我第一版选择分开纯粹为了逐算子验证时能精确定位误差来源。等所有算子都跑通了再做一个rms_norm_residual_fused作为第 22 号优化版这是后话。对于初学者来说先分后合永远比先合后分好。2.2 哪些算子容易低估RoPE 和 KV Cache 更新比想象中烦很多人一听到“手写 21 个算子”第一反应是难在矩阵乘。但等我实际写完才发现矩阵乘类算子反而是最省心的因为 Triton 的tl.dot已经封装得很好。真正容易翻车的是 RoPE 这种看似简单的 elementwise 算子和 KV Cache 更新这种跟内存布局强耦合的算子。RoPE 的麻烦在于它有两种实现方式连续的 cos/sin 预计算表以及按位置动态计算。我第一版直接按位置动态算每生成一个 token 都要算一组旋转角开销不大但 FP32 与 FP16 混用时会引入微小误差累积几十层后最终的 logits 跟 PyTorch 的 golden output 差异会达到 1e-2 量级。后来我改成预计算一张cos_cache和sin_cache按 position 索引误差立刻降到 1e-4 以下。这个改动本身很简单但它提示了一个原则所有跟位置相关的算子优先查表而不是临时算。KV Cache 更新则暴露了内存布局的问题。PyTorch 里 KV Cache 常用[layers, batch, kv_heads, seq_len, head_dim]的布局方便索引但不方便 kernel。我在 Triton kernel 里把它当作一个扁平的指针偏移来写每次更新时算好 offset。这个阶段花的时间比预想多因为cache_mask算子需要同时处理“当前序列长度”和“哪些位置是 padding”两件事如果一个 token 的 mask 没对齐后面 qk_dot 的 softmax 会直接把整个概率分布拉偏。2.3 为什么不用 torch.compile我需要的是“看得见”的算子这应该是很多人会问的问题PyTorch 2.x 的torch.compile不也能生成 Triton kernel 吗为什么还要手写我的回答是torch.compile生成的是给机器看的代码不是给人看的。它的 fusion 策略、block 切分、reduce 方式都是黑盒你无法在某一步出错时定位到具体算子。手写的意义在于每个 kernel 都能独立运行、独立对比数值、独立调参。这就像一个厨师如果用自动炒菜机可以做菜但想真正理解火候还是得自己颠几勺锅。所以这篇系列文章里我不会讨论怎么用编译器自动生成而是明确站在“手写算子做减法”的一侧。CUDAGraph 解决的是“kernel launch 太多”的 CPU 开销问题split-K 解决的是“单个 kernel 并行度不够”的 GPU 利用率问题这两者加上手写算子对内存读写的精细控制才是这条路线的真正价值。3. split-K 的两种打开方式GEMV 和 RMS Norm 的并行化split-K 这个术语听起来高大上本质就是一句话当一个归约维K 维太大、而输出维M 或 N 维太小的时候把 K 维切成好几段让多个 block 并行做部分归约最后再合并。这就像修一条很长的路如果只分给一个人从头走到尾耗时极长但如果把路分成几十段每段派一个人同时修最后把各段的成果拼接起来效率就上去了。3.1 GEMV 为什么要 split-K单 block 串行循环是性能黑洞生成阶段的矩阵乘全是 GEMV也就是M1的矩阵乘。拿q_proj来说输入是[1, 896]的隐藏状态权重是[num_q_heads * head_dim, hidden_size] [1024, 896]。如果按常规矩阵乘的写法输出 N 有 1024 个元素但每个输出元素都需要遍历 K896 维做累加。一个 block 处理一行输出的话内部是一个循环 896 次的串行累加GPU 的几千个核心只有一个在干活内存带宽跑不满延迟全堵在串行链路上。split-K 的做法是把 K 维切分。比如切成 4 段每个 block 只处理[1024, 224]的局部矩阵乘得到部分和最后用一个splitk_reducekernel 把这 4 份部分和加起来。这样并行度从“1024 个 block”变成“4096 个 block”每个 block 内部串行循环长度从 896 降到 224GPU 利用率立刻上来了。我当时实测的结果是不 split 时 q_proj 的耗时约 85 微秒切 4 段后降到 31 微秒切 8 段后降到 22 微秒。但注意不是越切越多越好——切到 16 段时耗时反而回升了因为归约开销超过了并行度收益。下面是一个参考实现的核心结构triton.jit def gemv_splitk_kernel( A_ptr, x_ptr, y_ptr, partial_ptr, M, N, K, stride_am, stride_ak, stride_xk, stride_yn, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, SPLIT_K: tl.constexpr, ): pid_n tl.program_id(0) pid_k tl.program_id(1) offs_n pid_n * BLOCK_N tl.arange(0, BLOCK_N) offs_k pid_k * BLOCK_K tl.arange(0, BLOCK_K) x_ptrs x_ptr offs_k x tl.load(x_ptrs, maskoffs_k K, other0.0) acc tl.zeros((BLOCK_N,), dtypetl.float32) for _ in range(0, BLOCK_K): a_ptrs A_ptr offs_n[:, None] * stride_am offs_k[None, :] * stride_ak a tl.load(a_ptrs, mask(offs_n[:, None] M) (offs_k[None, :] K), other0.0) acc tl.sum(a * x[None, :], axis1) offs_k BLOCK_K x_ptrs BLOCK_K x tl.load(x_ptrs, maskoffs_k K, other0.0) y_ptrs partial_ptr (pid_n * SPLIT_K pid_k) * BLOCK_N tl.store(y_ptrs, acc, maskoffs_n M)注意这里我用的是tl.sum(a * x)而不是tl.dot因为BLOCK_N和BLOCK_K都是普通数时小规模 GEMV 用广播乘法加归约比tl.dot更灵活尤其是当BLOCK_N不是 16 的倍数时。真正要提速可以把内层循环展开让一次迭代尽量多读几个 K 块。3.2 RMS Norm 的 split-K求平方和才是真正的归约瓶颈RMS Norm 比 GEMV 简单但同样有归约瓶颈。它的计算公式是y x / sqrt(mean(x^2) eps) * gamma其中核心是求sum(x^2)这是对整个 hidden_size896 维度的归约。一个 block 串行读完 896 个元素再算平方和单 block 延迟高尤其是在小 batch 场景下整个 GPU 只有少数几个 block 在工作。我的做法是让SPLIT_K个 block 并行读输入的不同部分每个 block 算局部的sum(x^2)写入一个[SPLIT_K]的临时数组然后由一个极小的 kernel 把这几个局部结果加起来reshape 成一个归一化系数再用 elementwise 算子把输入乘上系数和 gamma。因为 896 很小切 4 段就够了切多了反而浪费。triton.jit def rms_norm_splitk_kernel( x_ptr, partial_ptr, eps, K, SPLIT_K: tl.constexpr, BLOCK_K: tl.constexpr, ): pid_k tl.program_id(0) offs_k pid_k * BLOCK_K tl.arange(0, BLOCK_K) x tl.load(x_ptr offs_k, maskoffs_k K, other0.0) local_sum tl.sum(x * x, axis0) tl.store(partial_ptr pid_k, local_sum)然后是合并 kernel它只有一个 block把SPLIT_K个数加起来得到全局sum_sq然后算rms_eps 1 / sqrt(sum_sq / K eps)最后再单开一个 elementwise kernel 去 scale 和乘 gamma。有些优化版本会把最后一步也融进归一化 kernel但第一版我保持拆开方便逐算子核对。3.3 split-K 归约的精度教训顺序不重要但中间结果精度很重要split-K 最容易出的问题就是精度。我第一次跑通时最终输出和 PyTorch 的差异到了 1e-1 量级完全不可用。查找过程很有意思单独看每个 kernel 的输出跟 PyTorch 中间结果对比都正常但层层叠加后误差被放大。最后用torch.autograd.gradcheck级别的逐层排查才发现问题出在 split-K 的归约精度上。原因很简单多个 block 的部分和累加顺序是不确定的如果中间用 FP16 保存部分和误差会被放大。比如某些 block 算到 1000另一个 block 算到 -999二者相加会损失精度。解决方案是部分和一律用 FP32 存储只有最后一层输出才 cast 回 FP16。这个改动看起来不起眼但最终误差从 1e-1 降到了 1e-4。另外splitk_reducekernel 在合并时也要用 FP32 累加不能用tl.sum的默认 reduction 而不管中间类型。还有一个细节atomic_add做归约虽然代码短但多次 run 的结果会有微小差异因为 atomic 操作的顺序不确定。如果你的场景需要可复现性建议用第 18 号splitk_reduce做两级归约而不是直接在部分和数组上做 atomic_add。可复现性在调试阶段太重要了不然你会分不清性能波动和逻辑错误。4. CUDA Graph 捕获期的三个坑静态池、shape 固定和流处理算子全部跑通之后下一个大问题是 launch 开销。生成阶段每生成一个 token 要跑 21 个 kernel加上中间的张量操作PyTorch 侧有近 60 次 kernel launch。每次 launch 的 CPU 开销约 5~10 微秒累计下来超过了 GPU 计算时间本身。CUDA Graph 可以把这一整串 kernel 捕获成一张图然后每次只需要一次 launch 就能回放这是生成性能从 80 tokens/s 跳到 223 tokens/s 的关键一步。但 CUDA Graph 的捕获期极其“娇气”我踩了三个大坑每一个都值得单拎出来说。4.1 坑一捕获期内不能做显存分配CUDA Graph 捕获期间所有显存分配操作都会导致捕获失败或者产生无法回放的内存地址。这个问题在第一次捕获时几乎必现因为 PyTorch 的显存分配器是惰性的——很多中间张量在第一次跑时才真正分配。解决方法是捕获前先“预热”一遍把可能用到的张量全部分配好捕获期间只做计算和 kernel launch。我的做法是在 capture 之前先用自己的 kernel 跑一次完整的 decode 流程用torch.empty预分配 KV Cache、中间激活、split-K partial buffer、输出 logits然后强制torch.cuda.synchronize()。之后所有算子都使用这些预分配 buffer不用 PyTorch 的 autograd 张量只操作裸指针和预分配的 tensor。# 捕获前准备 kv_cache torch.empty(total_layers, 1, num_kv_heads, max_seq_len, head_dim, dtypetorch.float16, devicecuda) logits torch.empty(1, vocab_size, dtypetorch.float16, devicecuda) partial_sum torch.empty(SPLIT_K, dtypetorch.float32, devicecuda) # 预热 run_all_kernels(input_ids_static, kv_cache, logits, partial_sum) torch.cuda.synchronize() # 捕获 g torch.cuda.CUDAGraph() with torch.cuda.graph(g): run_all_kernels(input_ids_static, kv_cache, logits, partial_sum)这里有个细节input_ids_static必须是固定地址的 tensor后续回放时用copy_把新的 token id 拷进去而不是重新传一个新的 tensor。4.2 坑二动态 shape 和动态循环次数会被“固化”CUDA Graph 捕获时会把 kernel 的 grid、block、循环次数都固化下来。如果模型里有依赖运行时数据的if分支或者基于seq_len的循环图会把这个分支条件也固化成捕获时的值之后回放时即使条件变了走的还是旧路径。这对 KV Cache 的 mask 和 attention 长度是致命的。解决思路是固定 shape用 mask 处理变长。我在缓存中预分配了max_seq_len2048但实际 mask 会屏蔽超出当前真实长度的位置。这样kernel 的循环次数始终是对最大长度展开回放时通过 mask 控制实际计算范围。代价是长尾算力浪费但生成场景下 greedy decode 的 seq_len 是单调增加的使用 mask 开销远小于重新捕获一张图。4.3 坑三多流处理和同步操作不能在捕获区内出现CUDA Graph 捕获要求整个 launch 序列在同一个流上且捕获期间不能有跨流同步比如调用torch.cuda.synchronize()或者stream.wait_stream()否则会报错。我的代码里没有显式使用多流但最初为了调试方便在 capture 块里加了一个torch.cuda.synchronize()直接导致捕获失败。正确做法是捕获区之前同步一次捕获区内完全依赖 stream 的顺序语义捕获区外部做同步。如果你需要并行执行多个模型分片或者数据加载记得把数据加载放在另一个流上在 capture 前完成。CUDA Graph 只适合“纯计算图”回放不适合跨流依赖。另外捕获期如果用到 cuBLAS/cuDNN 的 workspace这些库的内部状态也可能被固化。如果 kernel 里调用了高级矩阵库最好给它们固定的 workspace 或者干脆像我一样全部用 Triton kernel避免黑盒状态。4.4 一个额外的提醒地址稳定性和冷热路径分离最后一个坑是我在对比中发现的CUDA Graph 回放的地址必须和捕获时完全一致。如果捕获后 KV Cache 被重新分配比如多跑了一次torch.empty之前捕获的图就会失效或者访问到旧内存。最简单的办法是把 KV Cache 和所有中间 buffer 放在一个模块里管理一次性分配之后整个生命周期不再改动。我实际还做了一个改进把“首次 token 的 prefill”和“后续 token 的 decode”分成两条路径。prefill 阶段输入序列长、shape 动态不适合 CUDA Graphdecode 阶段固定 batch1、固定 max_seq_len直接用图回放。两者共用一个 KV Cache在 prefill 结束后记录current_seq_lendecode 时用 mask 把前面已经填充的 KV 位置都计算进去。这样既保住了动态 prefill 的灵活性又享受了 decode 阶段 graph 的高吞吐。5. 调试实录与性能结果从 NaN 满天飞到 223 tokens/s调试手写 kernel 的过程本质上是“用 golden reference 逐算子对账”。我的 golden reference 不是整模型输出而是每层的中间结果。具体做法是先用 FP16 的 PyTorch 模型跑一遍把每个算子的输入输出都 dump 成 numpy 文件然后手写 kernel 跑同一份输入对比每个输出的 shape、dtype、数值分布。任何一层的误差超过阈值就锁定到那一个 kernel 里。5.1 三个真实的 debug 案例mask 值、归约顺序、地址偏移第一个大 bug 是 NaN 满天飞。现象是前几层输出正常到第 5 层之后 softmax 的结果出现 NaN。一开始我怀疑是 FP16 溢出反复检查数据范围都没问题。后来单步跑softmax_maskkernel发现 causal mask 的位置我用0去填充而不是-inf。虽然 mask 看起来只是“不参与计算”但 Triton 的tl.where在 mask 为0时会保留原始 logits 值导致 Attention Score 被放大到几千softmax 分母溢出变成 NaN。修正为-10000.0而不是0.0问题消失。这个教训是mask 要体现在数值上而不是逻辑上。第二个 bug 是 split-K 的归约结果偶尔出现“毛刺”。概率性复现不是每次都有。逐行打印部分和后发现q_proj的第 3 个 block 偶尔读到未初始化的显存。原因是splitk_reducekernel 里我用了固定长度的BLOCK_SPLIT_INDEX而实际的 SPLIT_K 小于这个值导致部分和数组尾部 4 个 element 没有初始化。修复方法是在splitk_reduce里也加 mask把超出 SPLIT_K 的部分忽略。第三个 bug 是地址偏移也是最难定位的一个。现象是生成的前 5 个 token 正常第 6 个 token 开始突然重复输出。查到最后发现是 KV Cache 的seq_len索引在kv_cache_update里没有加cache_mask提供的偏移导致新写入的位置恰好覆盖了旧内容。这个 bug 纯靠代码审查很难看出来最后是给每个算子加了一个“data stamp”日志打印它读写的 KV Cache 地址范围才在cache_update和qk_dot之间发现重叠。5.2 性能拆分到底是谁在拖慢生成速度当你把算子全手写之后性能分析就变得非常透明因为每个算子耗时可以直接测量不需要像黑盒 profiling 那样猜。我用torch.cuda.Event包住每个 kernel统计生成 100 个 token 的平均耗时分布。算子/环节平均耗时微秒占比备注q/k/v_proj GEMV38.212.1%split-K8qk_dot softmax pv_dot52.616.7%三个小 kernel 合并耗时attn_out_proj14.14.5%gate/up/down_proj61.419.5%MLP 三个 GEMV其余 elementwise/norm/mask42.313.4%CUDA Graph 回放总耗时315.7100%含 launch 开销这个拆分告诉我一个反直觉的事实手写算子的计算耗时只占一多半剩下的开销全在未融合的 elementwise 算子和 memory bound kernel 上。所以下一步提升的方向不是继续优化矩阵乘而是把residual_addrms_norm、silu_muldown_proj这类算子合并减少 kernel 数量。这也解释了为什么很多生产级推理框架会做“算子融合”而不是单纯压单个算子。5.3 最终结果与运行体验数据都在这里调试完成后用 greedy decoding 跑了一段约 256 token 的文本生成最终数据如下首 token 延迟prefill 阶段用的是普通 eager 模式没有进 CUDA Graph首 token 约 28ms主要开销在长序列 attention 的 mask 计算上。稳态生成速度CUDA Graph 回放后平均每 token 约 4.48ms折合 223 tokens/s。相比纯 PyTorch eager 模式84 tokens/s提升了 2.65 倍。显存占用静态 KV Cache 加上预分配中间 buffer总计约 2.8GB比 PyTorch 动态分配还低一点因为避免了中间张量的反复创建和释放。数值一致性手写 kernel 与 PyTorch FP16 模型最终输出的 token 序列完全一致greedy 场景下。这个体验完整复现了从“算子地图”到“性能数据”的闭环。整个过程最费时间的不是写 kernel而是调试和验证。6. 一些我写完后才真正理解的“为什么”文章最后分享几个我这次实践后才真正想明白的道理严格来说不算总结只是这些经验在动手前没人告诉我我觉得值得记下来。第一Triton kernel 的调参顺序比算法更影响性能。很多人一上来就追求 BLOCK_SIZE 最大化但实际上对 memory bound 的算子来说num_warps、num_stages、maxnreg这些编译期参数往往更关键。比如rms_norm这个算子固定 BLOCK_K896 时num_warps4比num_warps8快 15%因为 warp 太多时线程间同步的开销超过了并行收益。这类参数没法从理论推导只有用triton.testing.do_bench一个个试。第二逐算子验证时要用“缩放输入”而不是“小模型”。很多人为了快速调试会把 hidden_size 减小、层数减少。这有风险因为小 shape 下 block 边界和 mask 的边界情况可能被掩盖而真实 shape 下暴露的恰恰是这些边界问题。我的做法是保持真实 shape但用随机生成的固定输入跑这样既能覆盖边界又能快速定位。第三CUDA Graph 的收益并非白拿它会让所有动态性都变成显式状态。如果你要支持 beam search、多样本 batch、动态长度图的固定 shape 会带来很大限制。我的选择是只在 greedy batch1 的 decode 路径上用 graph其余路径继续 eager。这看起来“不酷”但工程上务实多了。第四也是我最想强调的手写算子的真正价值是让性能问题变得可见、可解释、可优化。用 PyTorch 时如果生成速度慢你只能看到“慢”这个结果用 Triton 手写后你能精确说出慢在哪个算子、慢在 memory bound 还是 compute bound、是该调 BLOCK 还是该融合算子。这个“可见性”才是性能工作的最底层能力。下篇内容我会拆这三个方向一是把residual_add rms_norm和silu_mul down_proj做融合目标是把 21 个算子合并到 14 个二是 prefill 阶段的tl.dot矩阵乘优化三是在 CUDA Graph 里处理 beam search 的工程方案。如果你也在折腾 Triton kernel欢迎收藏这篇地图等你的 NaN 报错了再回来看一眼第 5 节大概率能少走半天弯路。
返回列表