Triton Language

一句话定位

Triton 是编写高性能 GPU kernel 的语言和编译器，使开发者能以高层次方式表达 block/tile 计算，再生成 GPU 代码。

系统位置

1
2
3
4
attention / GEMM algorithm
→ PyTorch operator or custom op
→ Triton / CUDA kernel
→ GPU execution

边界与易混淆项

  • Triton Language：GPU kernel 编程与编译。
  • NVIDIA Triton Inference Server：模型服务系统。
  • TorchDynamo：torch.compile 图捕获组件。
  • NVIDIA Dynamo：分布式推理服务框架。

它们名称相似，但职责完全不同。

Kernel 优化的定位方法

1
2
3
4
5
6
端到端慢
→ profiler 找到热点 op
→ 判断算力受限还是带宽/访存受限
→ 设计 tile、布局、融合与 mask
→ Triton kernel 与 reference 数值对照
→ 在目标 shape、dtype、硬件上回归

Triton 只改变算子执行层，不能修复错误的 batching、KV 管理或网络路由。一个 kernel 在单个 shape 上变快，也可能因编译开销、寄存器压力或泛化 shape 变慢；因此报告应包含 warmup、编译时间、p50/p99、数值误差和实际工作负载。

相关节点：SDPA、AI System Performance Bottlenecks。

Triton 优化实际上在改什么

Triton kernel 的核心不是把 Python 代码“翻译成 GPU 代码”，而是显式选择数据如何按 tile 载入、在片上缓存中复用、写回以及如何处理边界 mask。一个 elementwise kernel 可能受 HBM 带宽限制；一个 GEMM/attention kernel 可能受矩阵维度、共享内存、寄存器或 occupancy 限制。BLOCK_SIZE、warp 数和 tile 布局是硬件相关的搜索空间，不是越大越好。

以 fused operation 为例，理论收益来自减少中间 tensor 的写回和 kernel launch：

1
2
未融合：load → op1 → global store → load → op2 → store
融合：  load → op1 → op2 → 一次 store

但融合会增大寄存器压力、降低可驻留 block 数，甚至让长尾 shape 退化。正确流程是对 reference 实现做逐元素/随机 shape 校验，再在目标 dtype、shape、batch 和硬件上测 warmup 后的 p50/p99；同时把首次编译时间和数值误差记录下来。

什么时候不该写 Triton

如果 profiler 显示主要时间在排队、KV 管理、NCCL、CPU tokenizer 或网络，重写 kernel 不会改变端到端瓶颈；如果目标 shape 很多且已有成熟 fused kernel，维护一个只在单一 shape 更快的 Triton 版本也可能增加回归成本。只有当热点算子稳定、reference 已验证、收益能覆盖编译和维护成本时，才值得把它纳入 production backend。

另外，Triton Language 与 NVIDIA Triton Inference Server 的名字相似但生命周期不同：前者产出 kernel，后者管理模型服务。前者的 benchmark 不能直接证明后者的吞吐或 SLO。