什么时候需要写 Triton 自定义算子？开发流程是什么？

先这样答

回答这个问题可以分为时机判断和开发流程两部分。在时机上，编写 Triton 自定义算子通常是最后的手段。只有当通过性能分析工具明确发现现有算子组合存在明显的瓶颈，例如访存浪费或中间结果频繁落地全局显存时，才需要考虑动手。常见的场景包括需要进行算子融合，比如将注意力的 Softmax 与掩码操作融合，或者将量化与反量化过程内联；另一种情况是模型引入了新算子，而现有的算子库中没有对应的实现。

Triton 本质上是一个使用 Python 编写 GPU kernel 的编译层。它采用了块级编程模型，编译器会自动处理底层的并行划分。这使得它的上手难度低于 CUDA C++，同时性能可以接近手写 CUDA，许多类似 FlashAttention 的工作都是以它为实现载体。

在实际开发流程中，第一步是确定分块策略，也就是规划 tile 的大小以及访存模式，这是决定算子性能的基础。第二步是编写 kernel，具体实现数据块的加载、计算以及最终结果的写回。第三步是正确性对拍，必须将自定义算子的输出与 PyTorch 的标准参考实现进行严格比对，确保计算逻辑无误。第四步是使用 autotune 进行自动调优，通过扫描不同的配置参数寻找当前硬件下的最优解。最后一步才是将其集成到实际的项目中。

在面试官面前可以这样收束：在面对性能优化需求时，大多数场景应该首先考虑使用 vLLM 或 PyTorch 中已经优化好的算子库，只有在现有库无法满足需求且性能瓶颈明确时，才启动 Triton 算子的开发。

面试官会怎么追问

• 「相比于写 CUDA C++，Triton 的块级编程模型有什么具体优势？」 传统 CUDA 编程需要开发者手动管理线程块、共享内存分配和显存合并访问。Triton 的块级编程模型将这些底层细节交由编译器自动处理，开发者只需要关注数据块层面的计算逻辑，大幅降低了开发门槛并提高了迭代效率。

• 「你刚才提到算子融合是为了减少中间结果落地，能具体解释一下过程吗？」 在未融合的情况下，多个连续的操作会将各自的输出写回全局显存，下一个操作再从显存读取，这会造成严重的访存开销。算子融合将这些操作放在一个 kernel 中执行，中间数据直接在寄存器或共享内存中传递，从而避免了多次读写全局显存。

• 「在对拍阶段如果发现精度对不上，你会从哪些方面排查？」 首先会检查数据类型转换是否一致，特别是混合精度计算时的截断问题。其次会检查分块计算的边界条件，确认掩码逻辑是否正确，防止在加载或写回数据时发生越界访问或错位。

回答的坑

• 认为只要为了性能优化就应该直接手写 Triton 算子。正确的逻辑是先做 profile 证明瓶颈所在，并且优先复用现有库，把自定义 kernel 作为最后的解决方案。
• 在描述开发流程时忽略了确定分块策略和自动调优环节。正确的答法需要体现出 Triton 开发中 tile 大小规划和 autotune 扫配置对最终性能的决定性作用。