花了两周时间把 Qwen3.5-0.8B 的 decoder 前向全部换成手写 Triton kernel 来跑。整个模型拆下来一共 21 个算子从 RMSNorm 到 QKV 投影、RoPE、flash attention 风格的 attention、MLP 三个 GEMM、最后的 lm_head全部用 Triton 重写了一遍最后用 CUDA Graph 把 decode 阶段的 kernel 启动封装成一次 launch。中间最费时间的三件事分别是算子对齐、split-K 的归约设计、以及 CUDA Graph 捕获时那一堆隐性问题。这篇文章是系列第一篇先给总体方案和这几个最关键的设计决策。后面我会挨个算子写清楚。如果你也想给一个小模型做纯手写 kernel 的推理框架这篇文章应该能帮你少走不少弯路。1. 为什么要把一个 0.8B 模型的所有算子都改成 Triton kernel先回答一个大家肯定会问的问题PyTorch 自带算子不香吗torch 里一个F.linear就搞定的事情为什么要自己写 kernel我的答案有三个分别对应 launch 开销、显存往返和算子融合。1.1 launch 开销在 decode 阶段比 kernel 本身还贵Qwen3.5-0.8B 这种小模型单层计算量其实很小。以 decode 阶段为例每生成一个 token要过 28 层 transformer每层又有 QKV 投影、attention、o_proj、MLP 三个 GEMM再加上 RMSNorm、残差、RoPE、KV cache 更新一个 token 的完整前向大概要发起上百次 kernel launch。每 launch 一次 CUDA kernelCPU 侧要准备参数、提交到 streamGPU 侧要调度、等前面的 kernel 结束这个固定开销大概是 5-10 微秒。如果 kernel 本身只跑 20 微秒那 launch 开销就占了 25%-30%。我在项目里试过纯 torch eager 模式跑 decode单 token 延迟 0.6ms 左右其中相当一部分不是算力不够而是 launch 太密。CUDA Graph 解决的就是这个问题把所有 kernel 提前捕获成一个图replay 的时候一次 launch 就能把整张图提交下去。手写 Triton kernel 加上 CUDA Graph等于把上百次 launch 压缩成一次。1.2 显存往返是带宽敏感型小模型的隐形杀手第二个收益来自算子融合。PyTorch 原生的算子链里很多中间结果要写回显存再用下一个算子读回来。比如 attention 的 score 矩阵torch 里是q k^T的结果先落地到 hbm然后 softmax 再从 hbm 读出来mask 也是一个单独的 tensor 要读一遍。对小模型来说SM 算力根本没有吃满瓶颈在显存带宽——每多一次往返就多花一次读和一次写的时间。Triton 手写 kernel 可以把这些操作融合在一个 kernel 内部完成。以我写的 attention kernel 为例qk_score、causal mask、softmax 全在同一个 kernel 里score 矩阵根本不出 SRAM。MLP 那边也是RoPE 直接在加载 q 和 k 的时候顺手用 fp32 算好乘进去不额外生成 cos/sin 的 bf16 tensor。1.3 为什么选 0.8B 而不是更大的模型0.8B 这个规模很适合做全手写 kernel 的验证。首先是单层计算量小算子本身的执行时间短launch 开销和显存带宽的影响会被放大手写 kernel 的优化效果更容易被看到。其次0.8B 的 config 里该有的机制一个不少GQA、RoPE、RMSNorm、tied embedding、bf16 训练权重全都覆盖到了。把这些机制对应的算子全部写一遍架构能力就基本建立起来了后面换 7B、14B 只是改 config 的问题。我当时还考虑过直接上 7B后来觉得小模型更容易定位问题算子对齐失败的时候单层代码就能复现不需要在 28 层里二分查 bug。事实证明这个决定是对的头几天我几乎每天都在改 kernel 里的 mask 或者 stride 错误如果换 7B调试成本会翻好几倍。1.4 prefill 和 decode 怎么共用同一套 kernel整个推理流程分两段prefill 处理输入序列shape 是(batch, seq_len)seq_len 可能是 256 也可能是 2048decode 逐 token 生成shape 是(batch, 1)刚启动时 1 个 token随着生成推进KV cache 的长度从 1 涨到几百甚至几千。我写的 21 个算子prefill 和 decode 共用同一份代码只是传入的 shape 不同。prefill 阶段走标准的 GEMM 和 flash attention 分支decode 阶段 M1走 split-K 分支和一些降维的 kernel 配置。这样维护成本低逻辑统一后面如果要加 beam search 或者 spec decode也只需要在同样的 kernel 上换参数。2. 从模型结构反推出来的 21 个算子清单动手写 kernel 之前得先把模型结构拆明白。我当时是直接读模型 config 和 safe tensor 的 shape 反推的这样最不容易出错。2.1 Qwen3.5-0.8B 的模型配置我本地用的这个 checkpointconfig 大概是这样的不同分支版本可能会差一点但结构一致配置项数值层数num_hidden_layers28隐藏层维度hidden_size1024query head 数8key/value head 数GQA2head_dim128intermediate_size4864词表大小151936最大序列长度32768这几个数字直接决定了算子怎么写。head_dim128意味着 attention 里的 qk 乘法 BLOCK 大小可以很规整GQA 的 kv head 只有 2 个decode 阶段需要广播到 8 个 query headintermediate_size4864这个数不是 16 的整数倍MLP 的 GEMM 里 BLOCK_N 的选择要特别小心不能默认用 4096 或 8192。2.2 21 个算子完整拆解表我把整个 decoder 前向拆成了下面这 21 个算子按执行顺序排序号算子名作用融合状态1rms_normattention 前的输入归一化独立 kernel2qkv_projx W_qkv^T输出 q/k/v 拼接fused GEMM3q_reshape_split把 q 按 8 个 head 拆开通过 stride 表达零成本4q_rope对 q 施加旋转位置编码独立 kernel5k_rope对 k 施加旋转位置编码独立 kernel6qk_scoreq k^T * scaleflash 风格7causal_mask_softmax下三角 mask softmax融合在 qk 后8pv_matmulp vflash 风格9attn_reshape_concat把 8 个 head 的 attention 输出合并回 hidden零成本10o_projattention 输出投影独立 GEMM11residual_addattention 后残差连接可与归一化融合这里独立12post_attn_rms_normMLP 前的归一化独立 kernel13mlp_gate_projMLP gate 分支 GEMM独立 GEMM14mlp_up_projMLP up 分支 GEMM独立 GEMM15act_mulSiLU(gate) * up融合 kernel16mlp_down_projMLP down 分支 GEMM独立 GEMM17residual_add_2MLP 后残差独立18kv_cache_store把新 k/v 写入 KV cache独立 kernel19cos_sin_precompute计算当前 token 的 RoPE 的 cos/sin独立 kernel20final_rms_norm最后一层输出归一化独立 kernel21lm_headembedding 转置 GEMM输出 logits独立 GEMM有人可能会问为什么把q_reshape_split和attn_reshape_concat单独列成算子它们不是零成本吗这是因为我需要一份清晰的算子清单来做调试和性能分析。在 Triton 里 view 和 transpose 确实不产生 kernel但我在写代码时把它们单独抽了一层负责把 QKV GEMM 的输出 strides 算对后面接入 flash attention 的时候才不会把维度搞混。这个习惯在写大模型推理时非常重要——维度问题是最隐蔽的一类 bug。2.3 哪些算子可以进一步融合我为什么没一开始就做从表格能看出来residual_add 和 rms_norm 其实可以融合成一个 kernelresidual 加完立刻做归一化省一次显存读写。我没在初版做融合是为了对齐方便。每个算子的 reference 都是最朴素的 torch 实现一旦融合reference 也要跟着改出 bug 时不好定位。我的原则是先保证每个算子独立正确再考虑融合优化。实际跑通之后等第二个版本再把residual_add rms_norm融合、把act_mul直接并进mlp_down_proj的加载阶段。初版最重要的是把算子和实际模型的对应关系理清楚这个对了后面怎么改都有保障。3. split-K 的设计decode 阶段真正需要它的是哪些算子split-K 是我这次写 GEMM 时花时间最多的一块也是最容易写错的地方。先说结论split-K 不是所有 GEMM 都需要甚至在很多场景用了反而更慢。3.1 为什么一个普通 GEMM 在 decode 场景会空转decode 阶段M1一个典型算子mlp_gate_proj的实际形状是(1, 1024) (1024, 4864)。如果用一个标准的 Triton matmul kernelBLOCK_M1, BLOCK_N128grid 大概是(1, 4864/128)也就是 38 个 block。你看 A100 有 108 个 SM这 38 个 block 分下去大部分 SM 是空闲的。更麻烦的是每个 block 内部要做一个 K1024 的循环。假设BLOCK_K256每个 block 串行循环 4 次每次加载A的一小块和B的一小块做点积。M 太小导致每个线程的复用率很低访存量依然不小但计算单元没有充分利用。split-K 的做法是把 K 维度也切成多块。还是这个mlp_gate_proj如果SPLIT_K4那 grid 变成(1, 38, 4)一共 152 个 block。每个 block 只负责 K256 的部分和并行度翻 4 倍每个 block 内部的循环次数也从 4 次降到了 1 次。代价是得到的是部分和需要额外一次归约。3.2 split-K 的 Triton kernel 实现我最初参考了 Triton 官方教程里 persistent matmul 的写法写了一个支持SPLIT_K参数的通用 GEMM。核心思路是这样的精简版重点看 split 部分triton.jit def matmul_kernel( A, B, C, M, N, K, stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, SPLIT_K: tl.constexpr, ): pid_m tl.program_id(0) pid_n tl.program_id(1) pid_k tl.program_id(2) offs_m pid_m * BLOCK_M tl.arange(0, BLOCK_M) offs_n pid_n * BLOCK_N tl.arange(0, BLOCK_N) # 每个 split 负责的 K 区间 split_k_dim K // SPLIT_K offs_k pid_k * split_k_dim tl.arange(0, BLOCK_K) a_ptrs A offs_m[:, None] * stride_am offs_k[None, :] * stride_ak b_ptrs B offs_k[:, None] * stride_bk offs_n[None, :] * stride_bn acc tl.zeros((BLOCK_M, BLOCK_N), dtypetl.float32) for _ in range(0, split_k_dim, BLOCK_K): a tl.load(a_ptrs, maskoffs_m[:, None] M, other0.0) b tl.load(b_ptrs, maskoffs_n[None, :] N, other0.0) acc tl.dot(a, b, acc) a_ptrs BLOCK_K * stride_ak b_ptrs BLOCK_K * stride_bk # 原子加或者部分和写出的逻辑在这里 out_ptrs C offs_m[:, None] * stride_cm offs_n[None, :] * stride_cn # 这里有两种模式 # 1) 全部原子加到 C 上避免二次 kernel # 2) 写到 (SPLIT_K, M, N) 的中间 buffer再单独 reduce关键是理解split_k_dim K // SPLIT_K和pid_k的关系每个 block 只循环自己负责的那一段 K最后得到一个部分和。这部分和的归约有两种常见做法下面单独说。3.3 原子加 vs 独立 reduce kernel我选了后者第一种做法是把 C 矩阵初始化为 0然后SPLIT_K个 block 做完部分和之后用tl.atomic_add累加。好处是省掉一个 kernel launch坏处是原子操作在同一块地址上竞争当 M 和 N 都很小、SPLIT_K又比较大时原子加的等待可能会抵消掉并行度提升。第二种做法是维护一个(SPLIT_K, M, N)的中间 buffer每个 block 把部分和写到自己的那一条最后用一个极小的 reduce kernel 把这SPLIT_K层加起来。多了一次 launch但每个 block 的写入完全并行没有竞争。我实际采用第二种。原因有三个一是在 CUDA Graph 捕获时原子操作虽然能工作但调试分析困难独立 reduce kernel 的性能行为一目了然二是中间 buffer 本身在 decode 阶段能被复用不会频繁分配三是后面要做 split-K 的误差分析和反量化时中间层级的分布信息都在排查问题方便。reduce kernel 很简单就是遍历SPLIT_K维做加法triton.jit def splitk_final_reduce_kernel( Partial, C, SPLIT_K: tl.constexpr, BLOCK_N: tl.constexpr, ): pid_m tl.program_id(0) pid_n tl.program_id(1) offs_n pid_n * BLOCK_N tl.arange(0, BLOCK_N) acc tl.zeros((BLOCK_N,), dtypetl.float32) for k in range(SPLIT_K): ptr Partial k * stride_partial_m pid_m * stride_partial_m offs_n acc tl.load(ptr) out_ptr C pid_m * stride_cm offs_n tl.store(out_ptr, acc)这个 kernel 的BLOCK_N我一般是 1024 或者 2048SPLIT_K通常取 2 或 4整个 kernel 运行时间在几微秒量级对整体延迟影响很小。3.4 哪些算子适合 split-K实测数据说话我把 21 个算子里的 GEMM 都测了一遍对比开启和不开 split-K 的效果A100-80Gbf16batch1decode 阶段算子MNK不开启SPLIT_K4结论qkv_proj1307210240.031ms0.038ms更慢不开o_proj1102410240.018ms0.025ms更慢不开mlp_gate_proj1486410240.042ms0.045ms基本打平mlp_up_proj1486410240.041ms0.044ms基本打平mlp_down_proj1102448640.036ms0.033ms略快可开lm_head115193610240.132ms0.096ms明显更快开有意思的是qkv_proj和o_proj用了 split-K 反而更慢。原因是它们的 K 都不大split-K 拆完之后每个 block 要算的部分和太少中间 buffer 的写入和 reduce kernel 的开销盖过了并行度收益。而lm_head的 N151936 非常大grid 的 N 方向本来就有很多 block但每个 block 的 M1 的 tile 太小加上 split-K 之后 K 维的并行度补上了综合下来提升约 27%。所以最后我在工程里只对lm_head和mlp_down_proj开了 split-K。这个选择不是拍脑袋是把每个 GEMM 形状跑一遍 benchmark 之后定下来的。3.5 split-K 的显存开销计算有一个容易忽略的细节split-K 需要一个(SPLIT_K, M, N)的中间 buffer。以lm_head为例M1, N151936, SPLIT_K4如果是 fp32 的 accumulator中间 buffer 是4 * 1 * 151936 * 4 bytes ≈ 2.4MB看起来不大。但如果你同时对mlp_down_proj也开 split-K再加上 attention 的临时 bufferdecode 阶段的临时显存就多了十几 MB。如果SPLIT_K取 8 或者 16中间 buffer 还会线性增长。这就是一个 trade-off并行度上去了显存占用量也上去了。我实测下来SPLIT_K4是收益和开销的平衡点SPLIT_K8对lm_head几乎没有进一步提升。4. CUDA Graph 捕获把 21 个 kernel 捏成一个 launch算子写完之后下一步就是把整个 decode 前向包进 CUDA Graph。这一步踩的坑比写算子还多我详细说说流程和陷阱。4.1 捕获前必须完成的 warmupCUDA Graph 捕获的原理是把捕获期间 stream 上提交的所有 kernel 记下来之后 replay 时不经过 CPU 侧的逐次 launch 逻辑。但这里有个坑Triton kernel 第一次加载时会有 lazy 编译和模块加载动作。具体来说Triton 的triton.jit函数在 Python 里第一次被调用时会触发 autotune 或者编译流程生成 cubin 并加载。这个加载过程涉及磁盘读取和驱动 API 调用如果发生在torch.cuda.graph捕获区间内轻则捕获失败重则捕获成功但 graph 里嵌套了隐式的模块加载调用replay 到第二次、第三次时就会出现不可预测的 crash。所以我的做法是在捕获之前先用同一个 shape 完整跑一遍整个 forward# warmup确保所有 Triton kernel 已完成编译和加载 for _ in range(3): logits model_forward(input_ids, kv_cache, position_ids) torch.cuda.synchronize()这里要跑 3 次第一次触发编译第二次确保所有 CUDA context 初始化和缓存分配都稳定下来第三次是保险。跑完synchronize()之后再进入 capture。另外所有需要在 graph 里复用的大 buffer 必须在 capture 之前分配好。KV cache 是一个典型的例子它在生成过程中会持续写入如果 capture 时才分配replay 时会复用同一块显存地址看起来没问题但实际上每次torch.empty在 capture 里都会被打进 graph 的依赖关系里极容易出问题。4.2 捕获代码的骨架代码结构大致是这样的self._graph torch.cuda.CUDAGraph() self._static_input_ids input_ids self._static_position_ids position_ids # warmup 之后正式开始捕获 torch.cuda.synchronize() with torch.cuda.graph(self._graph): self._static_logits model_forward( self._static_input_ids, self._kv_cache, self._static_position_ids, ) torch.cuda.synchronize()捕获完成后每次 decode 只需要self._static_input_ids.copy_(input_ids) self._static_position_ids.copy_(position_ids) self._graph.replay() logits self._static_logits注意输入是通过 copy 塞进静态 tensor 的输出也是从静态 tensor 里取。整个过程只有一个replay()调用CPU 侧几乎不花时间。4.3 捕获期间绝对不要出现这几类操作我在调试过程中踩了几个印象深刻的坑列出来给大家排雷第一torch.empty和动态分配。如果你的 forward 里用了torch.empty或者依赖每次 shape 重新分配的临时 tensor在 capture 阶段它会在 graph 里创建一个固定显存段的依赖。第一次 replay 可能正常第二次 replay 那部分数据已经被上一次覆盖结果就是数值错误而且是那种很难复现的间歇性错误。解决办法是提前分配好所有临时 buffer或者把临时 tensor 的分配移到 capture 外面。第二GPU 到 CPU 的同步操作。比如.item()、.cpu()、torch.cuda.synchronize()以及任何依赖 GPU 结果做 Python 分支判断的代码。在 capture 期间kernel 根本不会真实执行它只是被记录。如果你在 capture 里写了一个if tensor.item() 0.5:这样的分支capture 的时候会拿第一次运行的值决定走哪条路后面 replay 永远走那条路行为就完全错了。第三Python 侧 shape 分支。同理如果你的 forward 里有if seq_len 1024:这种条件capture 时只记录一次。decode 阶段 seq_len 每次都在变化如果走了没被记录的分支结果是未定义的。所以我让所有算子都接受固定 shapeseq_len 的变化靠 grid 和 mask 内部处理不改变 Python 控制流。第四Triton 的 autotune。我把所有算子都关闭了autotune手动指定 config。不是因为 autotune 不好而是它在 graph 捕获阶段会在后台跑 benchmark跟 capture 的确定性要求冲突。上线前可以 autotune 一把把最优 config 写死capture 时不再触发。4.4 为什么 decode 阶段受益最大prefill 为什么不用 graphdecode 是一个循环每生成一个 token跑一次前向采样把新的 token 拼回去继续下一次。这个循环要跑几百次每次前向里有几十个 kernel launch。CUDA Graph 一次 replay 替代了这几十次 launch收益是乘法级别的。prefill 阶段就不同了。prefill 只跑一次处理完整个输入序列一次性生成最后一个位置的 logits。它不在乎 launch 开销因为总时间主要被 GEMM 的执行时间占掉。而且 prefill 的 seq_len 每次都变如果强行 capture graph要么按长度分 bucket浪费显存要么每次重新 capturecapture 本身有开销得不偿失。所以我的设计是prefill 走普通 Triton kernel launchdecode 走 CUDA Graph。两套代码共用同样的 21 个算子只是在 decode 入口包了一层 graph。4.5 decode 循环的完整结构实际生成循环长这样input_ids torch.full((1, 1), bos_token_id, devicecuda, dtypetorch.long) position_ids torch.zeros((1, 1), devicecuda, dtypetorch.long) # prefill用传入的 prompt 先跑一次普通 kernel if prompt_ids.shape[1] 1: logits model_forward(prompt_ids, kv_cache, position_ids) # 采样得到第一个新 token 并更新 position # decode进入 graph replay 循环 for _ in range(max_new_tokens): self._static_input_ids.copy_(input_ids) self._static_position_ids.copy_(position_ids) self._graph.replay() logits self._static_logits next_token logits.argmax(dim-1) # 在 graph 外采样 input_ids next_token.reshape(1, 1) position_ids 1采样放到 graph 外面因为argmax返回的是 GPU tensor不涉及同步但在 CPU 侧获取 token id 时要注意是否真的需要同步。5. 算子对齐、精度边界与最终效果手写 kernel 的痛点是写出来跑得快不够结果必须正确。我花在对齐上的时间比写 kernel 本身还多。5.1 每个算子的 reference 怎么选我的对齐策略是每个算子都对着一个最简单的 torch 实现逐层验证而不是等到 21 个算子全部写完再端到端对比。单个算子出了问题单层就能复现不需要在 28 层里猜。以几个典型算子为例算子reference 实现rms_normtorch 手写 RMSNorm 公式qkv_projtorch.nn.functional.linear(x, qkv_weight)q_rope / k_rope先用 PyTorch 的 RoPE 实现生成 cos/sin再手动旋转qk_scoreq k.transpose(-2, -1) * scalecausal_mask_softmaxmasked_fillsoftmaxpv_matmulattn_probs vmlp 系列F.linear(x, gate_weight),F.linear(x, up_weight),F.linear(hidden, down_weight)每个算子对齐时我都写了一个断言函数比较 Triton 输出和 reference 的max_abs_diff和mean_abs_diff。如果 diff 超过了预设阈值就打印出算子在哪个 shape、哪个位置出现了最大偏差方便定位。5.2 bf16 下的误差容忍度0.8B 模型加载的是 bf16 权重所以在 Triton kernel 里输入也是 bf16。这里要注意几个误差来源。第一个是tl.dot的累加顺序。Triton 的 fp32 accumulator 是官方推荐的这点必须做不然误差会大很多。但即使有 fp32 accumulator不同 K 维度的累加顺序依然会影响最终数值。bf16 只有 8 位尾数当 K 很大时最后几位有效数字的差异是正常的。我的经验是GEMM 类算子max_abs_diff在 bf16 下允许到2e-2attention 的 logits 在1e-2量级完全达标。第二个是 flash attention 的在线 softmax 和 torch 的普通 softmax 数值不同。在线 softmax 是逐步更新的和一次性算完整 softmax 在数学上等价但浮点上不完全一样。好在 0.8B 模型深度有限这个差异在后续层会被压缩最终输出 logits 的 argmax 仍然和 torch reference 完全一致。第三个是 RoPE 的精度。cos/sin 表我在 Triton kernel 内部用 fp32 计算再乘到 bf16 的 q 和 k 上。torch 的 RoPE 实现如果是 float32 算的对比时要统一成同一种 dtype否则差异可能来自 RoPE 本身而不是 kernel 写错。5.3 最终的精度验证方法所有 21 个算子逐一对齐通过后我做了一次端到端验证拿同一个 prompt分别用 torch eager 模式和 Triton kernel CUDA Graph 模式跑一遍比较每条 token 的 logits 分布。结论是贪心解码输出的 token 序列完全一致top-kk50采样的候选集合也保持一致。更具体地说我用了一个 128 token 的 prompt跑了两组输出逐 token 对比 logits 的最大差值整体在3e-2以内算是相当干净了。这个结果验证了前面的判断bf16 下手写 kernel 完全能对齐 torch reference。5.4 实测性能数据最后给一组相对数据环境是 A100-80Gbf16batch1输入 prompt 256 tokens输出 128 tokens阶段torch eager手写 Triton无 graph手写 Triton split-K CUDA Graphprefill 256 tokens8.4ms6.1ms6.1msdecode 单 token0.62ms0.50ms0.44msdecode 吞吐估算~1600 token/s~2000 token/s~2270 token/s这个数字不是绝对的换卡、换驱动、换 CUDA 版本都会有波动。但从我的测试看手写 Triton 带来的提升主要来自算子融合和减少显存往返CUDA Graph 带来的提升主要来自消除 launch 开销split-K 则只对那几个指定的 GEMM 有帮助。三者叠加decode 端到端提升了约 40%。对比结果还有一点值得说prefill阶段 CUDA Graph 几乎没有帮助因为 prefill 本来就是一次性的launch 开销占比小。这验证了我前面的判断——CUDA Graph 要在循环次数多、launch 密集的场景才划算。最后分享一个我在这个项目里悟到的经验手写 21 个算子最大的收获不是那 40% 的性能提升而是理解了 kernel 从 launch 到执行到底经历了什么。以前我用 torch 的F.linear觉得那就是一次矩阵乘法写完 Triton GEMM 之后才意识到一个 kernel 的 grid 设计、block 大小、访存模式对最终耗时的影响有多大。如果你想复现这个项目我的建议是先从 RMSNorm 和 MLP 开始练手。这两个算子结构简单、对齐容易能让你快速建立对 Triton 的信心。然后再碰 attention最后再上 CUDA Graph。如果你也被某个算子的对齐问题卡住最有效的排查方法是把 shape 缩小到最小可复现的规模比如M2, K32, N64然后逐项打印 stride 和 mask——80% 的问题出在那里而不是数学公式。
手写Triton kernel重构Qwen3.5-0.8B：21算子全流程优化与CUDA Graph实践
花了两周时间把 Qwen3.5-0.8B 的 decoder 前向全部换成手写 Triton kernel 来跑。整个模型拆下来一共 21 个算子从 RMSNorm 到 QKV 投影、RoPE、flash attention 风格的 attention、MLP 三个 GEMM、最后的 lm_head全部用 Triton 重写了一遍最后用 CUDA Graph 把 decode 阶段的 kernel 启动封装成一次 launch。中间最费时间的三件事分别是算子对齐、split-K 的归约设计、以及 CUDA Graph 捕获时那一堆隐性问题。这篇文章是系列第一篇先给总体方案和这几个最关键的设计决策。后面我会挨个算子写清楚。如果你也想给一个小模型做纯手写 kernel 的推理框架这篇文章应该能帮你少走不少弯路。1. 为什么要把一个 0.8B 模型的所有算子都改成 Triton kernel先回答一个大家肯定会问的问题PyTorch 自带算子不香吗torch 里一个F.linear就搞定的事情为什么要自己写 kernel我的答案有三个分别对应 launch 开销、显存往返和算子融合。1.1 launch 开销在 decode 阶段比 kernel 本身还贵Qwen3.5-0.8B 这种小模型单层计算量其实很小。以 decode 阶段为例每生成一个 token要过 28 层 transformer每层又有 QKV 投影、attention、o_proj、MLP 三个 GEMM再加上 RMSNorm、残差、RoPE、KV cache 更新一个 token 的完整前向大概要发起上百次 kernel launch。每 launch 一次 CUDA kernelCPU 侧要准备参数、提交到 streamGPU 侧要调度、等前面的 kernel 结束这个固定开销大概是 5-10 微秒。如果 kernel 本身只跑 20 微秒那 launch 开销就占了 25%-30%。我在项目里试过纯 torch eager 模式跑 decode单 token 延迟 0.6ms 左右其中相当一部分不是算力不够而是 launch 太密。CUDA Graph 解决的就是这个问题把所有 kernel 提前捕获成一个图replay 的时候一次 launch 就能把整张图提交下去。手写 Triton kernel 加上 CUDA Graph等于把上百次 launch 压缩成一次。1.2 显存往返是带宽敏感型小模型的隐形杀手第二个收益来自算子融合。PyTorch 原生的算子链里很多中间结果要写回显存再用下一个算子读回来。比如 attention 的 score 矩阵torch 里是q k^T的结果先落地到 hbm然后 softmax 再从 hbm 读出来mask 也是一个单独的 tensor 要读一遍。对小模型来说SM 算力根本没有吃满瓶颈在显存带宽——每多一次往返就多花一次读和一次写的时间。Triton 手写 kernel 可以把这些操作融合在一个 kernel 内部完成。以我写的 attention kernel 为例qk_score、causal mask、softmax 全在同一个 kernel 里score 矩阵根本不出 SRAM。MLP 那边也是RoPE 直接在加载 q 和 k 的时候顺手用 fp32 算好乘进去不额外生成 cos/sin 的 bf16 tensor。1.3 为什么选 0.8B 而不是更大的模型0.8B 这个规模很适合做全手写 kernel 的验证。首先是单层计算量小算子本身的执行时间短launch 开销和显存带宽的影响会被放大手写 kernel 的优化效果更容易被看到。其次0.8B 的 config 里该有的机制一个不少GQA、RoPE、RMSNorm、tied embedding、bf16 训练权重全都覆盖到了。把这些机制对应的算子全部写一遍架构能力就基本建立起来了后面换 7B、14B 只是改 config 的问题。我当时还考虑过直接上 7B后来觉得小模型更容易定位问题算子对齐失败的时候单层代码就能复现不需要在 28 层里二分查 bug。事实证明这个决定是对的头几天我几乎每天都在改 kernel 里的 mask 或者 stride 错误如果换 7B调试成本会翻好几倍。1.4 prefill 和 decode 怎么共用同一套 kernel整个推理流程分两段prefill 处理输入序列shape 是(batch, seq_len)seq_len 可能是 256 也可能是 2048decode 逐 token 生成shape 是(batch, 1)刚启动时 1 个 token随着生成推进KV cache 的长度从 1 涨到几百甚至几千。我写的 21 个算子prefill 和 decode 共用同一份代码只是传入的 shape 不同。prefill 阶段走标准的 GEMM 和 flash attention 分支decode 阶段 M1走 split-K 分支和一些降维的 kernel 配置。这样维护成本低逻辑统一后面如果要加 beam search 或者 spec decode也只需要在同样的 kernel 上换参数。2. 从模型结构反推出来的 21 个算子清单动手写 kernel 之前得先把模型结构拆明白。我当时是直接读模型 config 和 safe tensor 的 shape 反推的这样最不容易出错。2.1 Qwen3.5-0.8B 的模型配置我本地用的这个 checkpointconfig 大概是这样的不同分支版本可能会差一点但结构一致配置项数值层数num_hidden_layers28隐藏层维度hidden_size1024query head 数8key/value head 数GQA2head_dim128intermediate_size4864词表大小151936最大序列长度32768这几个数字直接决定了算子怎么写。head_dim128意味着 attention 里的 qk 乘法 BLOCK 大小可以很规整GQA 的 kv head 只有 2 个decode 阶段需要广播到 8 个 query headintermediate_size4864这个数不是 16 的整数倍MLP 的 GEMM 里 BLOCK_N 的选择要特别小心不能默认用 4096 或 8192。2.2 21 个算子完整拆解表我把整个 decoder 前向拆成了下面这 21 个算子按执行顺序排序号算子名作用融合状态1rms_normattention 前的输入归一化独立 kernel2qkv_projx W_qkv^T输出 q/k/v 拼接fused GEMM3q_reshape_split把 q 按 8 个 head 拆开通过 stride 表达零成本4q_rope对 q 施加旋转位置编码独立 kernel5k_rope对 k 施加旋转位置编码独立 kernel6qk_scoreq k^T * scaleflash 风格7causal_mask_softmax下三角 mask softmax融合在 qk 后8pv_matmulp vflash 风格9attn_reshape_concat把 8 个 head 的 attention 输出合并回 hidden零成本10o_projattention 输出投影独立 GEMM11residual_addattention 后残差连接可与归一化融合这里独立12post_attn_rms_normMLP 前的归一化独立 kernel13mlp_gate_projMLP gate 分支 GEMM独立 GEMM14mlp_up_projMLP up 分支 GEMM独立 GEMM15act_mulSiLU(gate) * up融合 kernel16mlp_down_projMLP down 分支 GEMM独立 GEMM17residual_add_2MLP 后残差独立18kv_cache_store把新 k/v 写入 KV cache独立 kernel19cos_sin_precompute计算当前 token 的 RoPE 的 cos/sin独立 kernel20final_rms_norm最后一层输出归一化独立 kernel21lm_headembedding 转置 GEMM输出 logits独立 GEMM有人可能会问为什么把q_reshape_split和attn_reshape_concat单独列成算子它们不是零成本吗这是因为我需要一份清晰的算子清单来做调试和性能分析。在 Triton 里 view 和 transpose 确实不产生 kernel但我在写代码时把它们单独抽了一层负责把 QKV GEMM 的输出 strides 算对后面接入 flash attention 的时候才不会把维度搞混。这个习惯在写大模型推理时非常重要——维度问题是最隐蔽的一类 bug。2.3 哪些算子可以进一步融合我为什么没一开始就做从表格能看出来residual_add 和 rms_norm 其实可以融合成一个 kernelresidual 加完立刻做归一化省一次显存读写。我没在初版做融合是为了对齐方便。每个算子的 reference 都是最朴素的 torch 实现一旦融合reference 也要跟着改出 bug 时不好定位。我的原则是先保证每个算子独立正确再考虑融合优化。实际跑通之后等第二个版本再把residual_add rms_norm融合、把act_mul直接并进mlp_down_proj的加载阶段。初版最重要的是把算子和实际模型的对应关系理清楚这个对了后面怎么改都有保障。3. split-K 的设计decode 阶段真正需要它的是哪些算子split-K 是我这次写 GEMM 时花时间最多的一块也是最容易写错的地方。先说结论split-K 不是所有 GEMM 都需要甚至在很多场景用了反而更慢。3.1 为什么一个普通 GEMM 在 decode 场景会空转decode 阶段M1一个典型算子mlp_gate_proj的实际形状是(1, 1024) (1024, 4864)。如果用一个标准的 Triton matmul kernelBLOCK_M1, BLOCK_N128grid 大概是(1, 4864/128)也就是 38 个 block。你看 A100 有 108 个 SM这 38 个 block 分下去大部分 SM 是空闲的。更麻烦的是每个 block 内部要做一个 K1024 的循环。假设BLOCK_K256每个 block 串行循环 4 次每次加载A的一小块和B的一小块做点积。M 太小导致每个线程的复用率很低访存量依然不小但计算单元没有充分利用。split-K 的做法是把 K 维度也切成多块。还是这个mlp_gate_proj如果SPLIT_K4那 grid 变成(1, 38, 4)一共 152 个 block。每个 block 只负责 K256 的部分和并行度翻 4 倍每个 block 内部的循环次数也从 4 次降到了 1 次。代价是得到的是部分和需要额外一次归约。3.2 split-K 的 Triton kernel 实现我最初参考了 Triton 官方教程里 persistent matmul 的写法写了一个支持SPLIT_K参数的通用 GEMM。核心思路是这样的精简版重点看 split 部分triton.jit def matmul_kernel( A, B, C, M, N, K, stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, SPLIT_K: tl.constexpr, ): pid_m tl.program_id(0) pid_n tl.program_id(1) pid_k tl.program_id(2) offs_m pid_m * BLOCK_M tl.arange(0, BLOCK_M) offs_n pid_n * BLOCK_N tl.arange(0, BLOCK_N) # 每个 split 负责的 K 区间 split_k_dim K // SPLIT_K offs_k pid_k * split_k_dim tl.arange(0, BLOCK_K) a_ptrs A offs_m[:, None] * stride_am offs_k[None, :] * stride_ak b_ptrs B offs_k[:, None] * stride_bk offs_n[None, :] * stride_bn acc tl.zeros((BLOCK_M, BLOCK_N), dtypetl.float32) for _ in range(0, split_k_dim, BLOCK_K): a tl.load(a_ptrs, maskoffs_m[:, None] M, other0.0) b tl.load(b_ptrs, maskoffs_n[None, :] N, other0.0) acc tl.dot(a, b, acc) a_ptrs BLOCK_K * stride_ak b_ptrs BLOCK_K * stride_bk # 原子加或者部分和写出的逻辑在这里 out_ptrs C offs_m[:, None] * stride_cm offs_n[None, :] * stride_cn # 这里有两种模式 # 1) 全部原子加到 C 上避免二次 kernel # 2) 写到 (SPLIT_K, M, N) 的中间 buffer再单独 reduce关键是理解split_k_dim K // SPLIT_K和pid_k的关系每个 block 只循环自己负责的那一段 K最后得到一个部分和。这部分和的归约有两种常见做法下面单独说。3.3 原子加 vs 独立 reduce kernel我选了后者第一种做法是把 C 矩阵初始化为 0然后SPLIT_K个 block 做完部分和之后用tl.atomic_add累加。好处是省掉一个 kernel launch坏处是原子操作在同一块地址上竞争当 M 和 N 都很小、SPLIT_K又比较大时原子加的等待可能会抵消掉并行度提升。第二种做法是维护一个(SPLIT_K, M, N)的中间 buffer每个 block 把部分和写到自己的那一条最后用一个极小的 reduce kernel 把这SPLIT_K层加起来。多了一次 launch但每个 block 的写入完全并行没有竞争。我实际采用第二种。原因有三个一是在 CUDA Graph 捕获时原子操作虽然能工作但调试分析困难独立 reduce kernel 的性能行为一目了然二是中间 buffer 本身在 decode 阶段能被复用不会频繁分配三是后面要做 split-K 的误差分析和反量化时中间层级的分布信息都在排查问题方便。reduce kernel 很简单就是遍历SPLIT_K维做加法triton.jit def splitk_final_reduce_kernel( Partial, C, SPLIT_K: tl.constexpr, BLOCK_N: tl.constexpr, ): pid_m tl.program_id(0) pid_n tl.program_id(1) offs_n pid_n * BLOCK_N tl.arange(0, BLOCK_N) acc tl.zeros((BLOCK_N,), dtypetl.float32) for k in range(SPLIT_K): ptr Partial k * stride_partial_m pid_m * stride_partial_m offs_n acc tl.load(ptr) out_ptr C pid_m * stride_cm offs_n tl.store(out_ptr, acc)这个 kernel 的BLOCK_N我一般是 1024 或者 2048SPLIT_K通常取 2 或 4整个 kernel 运行时间在几微秒量级对整体延迟影响很小。3.4 哪些算子适合 split-K实测数据说话我把 21 个算子里的 GEMM 都测了一遍对比开启和不开 split-K 的效果A100-80Gbf16batch1decode 阶段算子MNK不开启SPLIT_K4结论qkv_proj1307210240.031ms0.038ms更慢不开o_proj1102410240.018ms0.025ms更慢不开mlp_gate_proj1486410240.042ms0.045ms基本打平mlp_up_proj1486410240.041ms0.044ms基本打平mlp_down_proj1102448640.036ms0.033ms略快可开lm_head115193610240.132ms0.096ms明显更快开有意思的是qkv_proj和o_proj用了 split-K 反而更慢。原因是它们的 K 都不大split-K 拆完之后每个 block 要算的部分和太少中间 buffer 的写入和 reduce kernel 的开销盖过了并行度收益。而lm_head的 N151936 非常大grid 的 N 方向本来就有很多 block但每个 block 的 M1 的 tile 太小加上 split-K 之后 K 维的并行度补上了综合下来提升约 27%。所以最后我在工程里只对lm_head和mlp_down_proj开了 split-K。这个选择不是拍脑袋是把每个 GEMM 形状跑一遍 benchmark 之后定下来的。3.5 split-K 的显存开销计算有一个容易忽略的细节split-K 需要一个(SPLIT_K, M, N)的中间 buffer。以lm_head为例M1, N151936, SPLIT_K4如果是 fp32 的 accumulator中间 buffer 是4 * 1 * 151936 * 4 bytes ≈ 2.4MB看起来不大。但如果你同时对mlp_down_proj也开 split-K再加上 attention 的临时 bufferdecode 阶段的临时显存就多了十几 MB。如果SPLIT_K取 8 或者 16中间 buffer 还会线性增长。这就是一个 trade-off并行度上去了显存占用量也上去了。我实测下来SPLIT_K4是收益和开销的平衡点SPLIT_K8对lm_head几乎没有进一步提升。4. CUDA Graph 捕获把 21 个 kernel 捏成一个 launch算子写完之后下一步就是把整个 decode 前向包进 CUDA Graph。这一步踩的坑比写算子还多我详细说说流程和陷阱。4.1 捕获前必须完成的 warmupCUDA Graph 捕获的原理是把捕获期间 stream 上提交的所有 kernel 记下来之后 replay 时不经过 CPU 侧的逐次 launch 逻辑。但这里有个坑Triton kernel 第一次加载时会有 lazy 编译和模块加载动作。具体来说Triton 的triton.jit函数在 Python 里第一次被调用时会触发 autotune 或者编译流程生成 cubin 并加载。这个加载过程涉及磁盘读取和驱动 API 调用如果发生在torch.cuda.graph捕获区间内轻则捕获失败重则捕获成功但 graph 里嵌套了隐式的模块加载调用replay 到第二次、第三次时就会出现不可预测的 crash。所以我的做法是在捕获之前先用同一个 shape 完整跑一遍整个 forward# warmup确保所有 Triton kernel 已完成编译和加载 for _ in range(3): logits model_forward(input_ids, kv_cache, position_ids) torch.cuda.synchronize()这里要跑 3 次第一次触发编译第二次确保所有 CUDA context 初始化和缓存分配都稳定下来第三次是保险。跑完synchronize()之后再进入 capture。另外所有需要在 graph 里复用的大 buffer 必须在 capture 之前分配好。KV cache 是一个典型的例子它在生成过程中会持续写入如果 capture 时才分配replay 时会复用同一块显存地址看起来没问题但实际上每次torch.empty在 capture 里都会被打进 graph 的依赖关系里极容易出问题。4.2 捕获代码的骨架代码结构大致是这样的self._graph torch.cuda.CUDAGraph() self._static_input_ids input_ids self._static_position_ids position_ids # warmup 之后正式开始捕获 torch.cuda.synchronize() with torch.cuda.graph(self._graph): self._static_logits model_forward( self._static_input_ids, self._kv_cache, self._static_position_ids, ) torch.cuda.synchronize()捕获完成后每次 decode 只需要self._static_input_ids.copy_(input_ids) self._static_position_ids.copy_(position_ids) self._graph.replay() logits self._static_logits注意输入是通过 copy 塞进静态 tensor 的输出也是从静态 tensor 里取。整个过程只有一个replay()调用CPU 侧几乎不花时间。4.3 捕获期间绝对不要出现这几类操作我在调试过程中踩了几个印象深刻的坑列出来给大家排雷第一torch.empty和动态分配。如果你的 forward 里用了torch.empty或者依赖每次 shape 重新分配的临时 tensor在 capture 阶段它会在 graph 里创建一个固定显存段的依赖。第一次 replay 可能正常第二次 replay 那部分数据已经被上一次覆盖结果就是数值错误而且是那种很难复现的间歇性错误。解决办法是提前分配好所有临时 buffer或者把临时 tensor 的分配移到 capture 外面。第二GPU 到 CPU 的同步操作。比如.item()、.cpu()、torch.cuda.synchronize()以及任何依赖 GPU 结果做 Python 分支判断的代码。在 capture 期间kernel 根本不会真实执行它只是被记录。如果你在 capture 里写了一个if tensor.item() 0.5:这样的分支capture 的时候会拿第一次运行的值决定走哪条路后面 replay 永远走那条路行为就完全错了。第三Python 侧 shape 分支。同理如果你的 forward 里有if seq_len 1024:这种条件capture 时只记录一次。decode 阶段 seq_len 每次都在变化如果走了没被记录的分支结果是未定义的。所以我让所有算子都接受固定 shapeseq_len 的变化靠 grid 和 mask 内部处理不改变 Python 控制流。第四Triton 的 autotune。我把所有算子都关闭了autotune手动指定 config。不是因为 autotune 不好而是它在 graph 捕获阶段会在后台跑 benchmark跟 capture 的确定性要求冲突。上线前可以 autotune 一把把最优 config 写死capture 时不再触发。4.4 为什么 decode 阶段受益最大prefill 为什么不用 graphdecode 是一个循环每生成一个 token跑一次前向采样把新的 token 拼回去继续下一次。这个循环要跑几百次每次前向里有几十个 kernel launch。CUDA Graph 一次 replay 替代了这几十次 launch收益是乘法级别的。prefill 阶段就不同了。prefill 只跑一次处理完整个输入序列一次性生成最后一个位置的 logits。它不在乎 launch 开销因为总时间主要被 GEMM 的执行时间占掉。而且 prefill 的 seq_len 每次都变如果强行 capture graph要么按长度分 bucket浪费显存要么每次重新 capturecapture 本身有开销得不偿失。所以我的设计是prefill 走普通 Triton kernel launchdecode 走 CUDA Graph。两套代码共用同样的 21 个算子只是在 decode 入口包了一层 graph。4.5 decode 循环的完整结构实际生成循环长这样input_ids torch.full((1, 1), bos_token_id, devicecuda, dtypetorch.long) position_ids torch.zeros((1, 1), devicecuda, dtypetorch.long) # prefill用传入的 prompt 先跑一次普通 kernel if prompt_ids.shape[1] 1: logits model_forward(prompt_ids, kv_cache, position_ids) # 采样得到第一个新 token 并更新 position # decode进入 graph replay 循环 for _ in range(max_new_tokens): self._static_input_ids.copy_(input_ids) self._static_position_ids.copy_(position_ids) self._graph.replay() logits self._static_logits next_token logits.argmax(dim-1) # 在 graph 外采样 input_ids next_token.reshape(1, 1) position_ids 1采样放到 graph 外面因为argmax返回的是 GPU tensor不涉及同步但在 CPU 侧获取 token id 时要注意是否真的需要同步。5. 算子对齐、精度边界与最终效果手写 kernel 的痛点是写出来跑得快不够结果必须正确。我花在对齐上的时间比写 kernel 本身还多。5.1 每个算子的 reference 怎么选我的对齐策略是每个算子都对着一个最简单的 torch 实现逐层验证而不是等到 21 个算子全部写完再端到端对比。单个算子出了问题单层就能复现不需要在 28 层里猜。以几个典型算子为例算子reference 实现rms_normtorch 手写 RMSNorm 公式qkv_projtorch.nn.functional.linear(x, qkv_weight)q_rope / k_rope先用 PyTorch 的 RoPE 实现生成 cos/sin再手动旋转qk_scoreq k.transpose(-2, -1) * scalecausal_mask_softmaxmasked_fillsoftmaxpv_matmulattn_probs vmlp 系列F.linear(x, gate_weight),F.linear(x, up_weight),F.linear(hidden, down_weight)每个算子对齐时我都写了一个断言函数比较 Triton 输出和 reference 的max_abs_diff和mean_abs_diff。如果 diff 超过了预设阈值就打印出算子在哪个 shape、哪个位置出现了最大偏差方便定位。5.2 bf16 下的误差容忍度0.8B 模型加载的是 bf16 权重所以在 Triton kernel 里输入也是 bf16。这里要注意几个误差来源。第一个是tl.dot的累加顺序。Triton 的 fp32 accumulator 是官方推荐的这点必须做不然误差会大很多。但即使有 fp32 accumulator不同 K 维度的累加顺序依然会影响最终数值。bf16 只有 8 位尾数当 K 很大时最后几位有效数字的差异是正常的。我的经验是GEMM 类算子max_abs_diff在 bf16 下允许到2e-2attention 的 logits 在1e-2量级完全达标。第二个是 flash attention 的在线 softmax 和 torch 的普通 softmax 数值不同。在线 softmax 是逐步更新的和一次性算完整 softmax 在数学上等价但浮点上不完全一样。好在 0.8B 模型深度有限这个差异在后续层会被压缩最终输出 logits 的 argmax 仍然和 torch reference 完全一致。第三个是 RoPE 的精度。cos/sin 表我在 Triton kernel 内部用 fp32 计算再乘到 bf16 的 q 和 k 上。torch 的 RoPE 实现如果是 float32 算的对比时要统一成同一种 dtype否则差异可能来自 RoPE 本身而不是 kernel 写错。5.3 最终的精度验证方法所有 21 个算子逐一对齐通过后我做了一次端到端验证拿同一个 prompt分别用 torch eager 模式和 Triton kernel CUDA Graph 模式跑一遍比较每条 token 的 logits 分布。结论是贪心解码输出的 token 序列完全一致top-kk50采样的候选集合也保持一致。更具体地说我用了一个 128 token 的 prompt跑了两组输出逐 token 对比 logits 的最大差值整体在3e-2以内算是相当干净了。这个结果验证了前面的判断bf16 下手写 kernel 完全能对齐 torch reference。5.4 实测性能数据最后给一组相对数据环境是 A100-80Gbf16batch1输入 prompt 256 tokens输出 128 tokens阶段torch eager手写 Triton无 graph手写 Triton split-K CUDA Graphprefill 256 tokens8.4ms6.1ms6.1msdecode 单 token0.62ms0.50ms0.44msdecode 吞吐估算~1600 token/s~2000 token/s~2270 token/s这个数字不是绝对的换卡、换驱动、换 CUDA 版本都会有波动。但从我的测试看手写 Triton 带来的提升主要来自算子融合和减少显存往返CUDA Graph 带来的提升主要来自消除 launch 开销split-K 则只对那几个指定的 GEMM 有帮助。三者叠加decode 端到端提升了约 40%。对比结果还有一点值得说prefill阶段 CUDA Graph 几乎没有帮助因为 prefill 本来就是一次性的launch 开销占比小。这验证了我前面的判断——CUDA Graph 要在循环次数多、launch 密集的场景才划算。最后分享一个我在这个项目里悟到的经验手写 21 个算子最大的收获不是那 40% 的性能提升而是理解了 kernel 从 launch 到执行到底经历了什么。以前我用 torch 的F.linear觉得那就是一次矩阵乘法写完 Triton GEMM 之后才意识到一个 kernel 的 grid 设计、block 大小、访存模式对最终耗时的影响有多大。如果你想复现这个项目我的建议是先从 RMSNorm 和 MLP 开始练手。这两个算子结构简单、对齐容易能让你快速建立对 Triton 的信心。然后再碰 attention最后再上 CUDA Graph。如果你也被某个算子的对齐问题卡住最有效的排查方法是把 shape 缩小到最小可复现的规模比如M2, K32, N64然后逐项打印 stride 和 mask——80% 的问题出在那里而不是数学公式。
