Zhao Dongyu's Blog

A life which is unexamined is not worth living.

0%

triton

Posted on Edited on In
Symbols count in article: 3.9k Reading time ≈ 4 mins.

其实最近心思一直在把之前的推理框架的工作能够沉淀一下发一篇论文，今天丁大佬教育我不要闭门造车，不要局限于推理框架，多看看别人的工程学习学习,学学triton、TVM、mlc-llm这些东西。很受用。

听人劝，吃饱饭。开始学习 triton

官网 翻译

Triton 是一种类似 Python 的开源编程语言，它使没有 CUDA 经验的研究人员能够编写高效的 GPU 代码——大多数时候与专家能够编写的代码相当。

其他框架too verbose, lack flexibility or generate code noticeably slower than Triton‘s hand-tuned baselines.

  • 目前已有的 DSL 在灵活性和（对相同算法）速度上明显慢于像 cuBLAS、cuDNN 或 TensorRT 这样的库中可用的最佳手写计算内核。
  • 已有的 DSL 如 polyhedral machinery (Tiramisu/Tensor Comprehensions)、scheduling languages (Halide、TVM) 等在效率上还有提升空间。

The architecture of modern GPUs can be roughly divided into three major components—DRAM, SRAM and ALUs—each of which must be considered when optimizing CUDA code:

  • Memory transfers from DRAM must be coalesced into large transactions to leverage the large bus width of modern memory interfaces.

  • Data must be manually stashed to SRAM prior to being re-used, and managed so as to minimize shared memory bank conflicts upon retrieval.

  • Computations must be partitioned and scheduled carefully, both across and within Streaming Multiprocessors (SMs), so as to promote instruction/thread-level parallelism and leverage special-purpose ALUs (e.g., tensor cores).

即

  • 合并访问 (Memory Coalescing)
  • 避免 Bank Conflict (Avoiding Shared Memory Bank Conflicts)
  • 提升计算效率 (Maximizing Computational Efficiency)

Triton 简化了 CUDA 编程的底层优化，通过自动化的方式减少开发者的负担，但同时保留了一些高层次的算法决策，例如 SM 任务调度和平铺策略，供开发者灵活调整。这种设计既提高了编程效率，又确保了对复杂算法的适用性。

Triton 通过对 blocks（尺寸为 2 的幂的小数组）的操作来公开实例内并行性，而不是单指令、多线程 (SIMT) 执行模型。通过这样做，Triton 有效地抽象了与 CUDA 线程块内的并发性相关的所有问题（例如，内存合并、共享内存同步/冲突、张量核心调度）。

在标准 CUDA 实现中，Softmax 归一化需要显式管理线程之间的同步（如多线程并行计算同一行的指数值和总和）。Triton 则通过将计算任务划分为每个 kernel 处理一行，避免了线程同步的复杂性，大大降低了代码实现难度，同时保持高性能。

The bottom line here is not that Triton is inherently better, but that it simplifies the development of specialized kernels that can be much faster than those found in general-purpose libraries.

The good performance of Triton comes from a modular system architecture centered around Triton-IR, an LLVM-based intermediate representation in which multi-dimensional blocks of values are first-class citizens.

总之， Triton 的核心理念是基于分块的编程范式可以促进神经网络的高性能计算核心的构建。CUDA 编写属于传统的 “单程序，多数据” GPU 执行模型，在线程的细粒度上进行编程，Triton 是在分块的细粒度上进行编程。

跟着教程走：https://triton-lang.org/main/getting-started/tutorials/index.html

OpenAI Triton 入门教程

triton 在循环中是逐块进行计算的。这种方法的一个关键优势是，它导致了块结构的迭代空间，相较于现有的DSL，为程序员在实现稀疏操作时提供了更多的灵活性，同时允许编译器为数据局部性和并行性进行积极的优化。

triton.jit

triton.jit 是 Triton 框架中的一个装饰器，用于将 Python 函数 JIT（即时）编译为高性能的 GPU 内核。

  1. JIT 编译：

    • 使用 Triton 编译器将被装饰的函数编译为 GPU 内核代码。
    • 编译后的内核在 GPU 上运行，利用 GPU 的并行计算能力实现高性能。

  2. 隐式转换：

    • 当调用一个使用 triton.jit 编译的函数时，传入的参数如果具有 .data_ptr() 方法（如 PyTorch 张量）和 .dtype 属性，会自动被转换为指针。
    • 这使得与 PyTorch 等框架的张量交互更加方便。

  3. 限制环境：

    • JIT 编译后的函数在 GPU 上运行，不能直接访问普通的 Python 对象或标准库。
    • 仅能使用以下资源：
    • Python 的基本类型（如 int、float）。
    • Triton 包中的内置函数。
    • 函数的参数。
    • 其他被 JIT 编译的函数。

triton.autotune

triton.autotune 是一个装饰器，用于为 triton.jit 编译的函数添加 自动调优 功能。它通过评估多个配置（configs），选择最优的内核配置以获得最佳性能。

核心功能

  1. 自动配置选择：
    • 提供多个候选配置（configs），triton.autotune 会自动评估这些配置的性能。
    • 适用于 GPU 上的高性能计算，能够动态调整块大小（如 BLOCK_SIZE）、线程数等关键参数。
  2. 动态适应输入：
    • 通过 key 参数指定，当输入参数发生变化时，重新评估配置的性能。
  3. 安全优化：
    • 通过 reset_to_zero 和 restore_value 确保多次内核调用不会造成副作用。
  4. 性能输出：
    • 如果设置环境变量 TRITON_PRINT_AUTOTUNING=1，Triton 会打印调优过程的日志，包括时间消耗和最佳配置。

triton.heuristics

triton.heuristics 是 Triton 提供的装饰器，用于动态计算 meta-parameter 的值。这对于以下情况非常有用：

  1. 自动调优代价高昂：在一些情况下，使用 triton.autotune 进行全面的自动调优可能会非常耗时。

  2. 自动调优不可用：某些情况下，输入特性使得自动调优不适用。

  3. 动态计算 meta-parameter：

    • 根据输入参数动态推导内核中的 meta-parameter。
    • 避免硬编码或繁琐的手动配置。
    • 例如：根据张量大小动态设置块大小 BLOCK_SIZE。

  4. 替代部分自动调优：

    • 相比 triton.autotune 的多次配置测试，triton.heuristics 通过计算直接生成最优或次优的 meta-parameter 值。

主要看一下 Matrix Multiplication 这个代码:

实现一个高性能的 FP16 矩阵乘法内核，其性能可以与 NVIDIA 的 cuBLAS 或 AMD 的 rocBLAS 相媲美。这些库是由硬件厂商提供的优化矩阵操作库，但它们通常是闭源的，不易定制。通过本教程，用户可以学习如何自行实现高性能矩阵乘法，同时能够轻松扩展或调整以适应特殊需求，例如深度学习中的自定义操作。

关键学习点

  1. 块级矩阵乘法（Block-level Matrix Multiplications）：
    • 通过将矩阵分成小块，逐块进行操作，以提升性能和内存利用率。
    • 这种分块操作可以显著提升 GPU 的计算效率。
  2. 多维指针运算（Multi-dimensional Pointer Arithmetic）：
    • 学习如何通过指针运算访问矩阵中的子块（block）。
    • 对于行优先（row-major）存储的矩阵，需要通过复杂的指针偏移计算每个块的位置。
  3. 程序重排序以提高 L2 缓存命中率（Program Re-ordering for Improved L2 Cache Hit Rate）：
    • 通过改变程序块的执行顺序（如按列优先或分组的方式执行），可以提高 L2 缓存的命中率。
    • 这种优化减少了从全局内存中加载数据的次数，从而显著提高性能。
  4. 自动性能调优（Automatic Performance Tuning）：
    • 利用 Triton 提供的自动调优工具，为特定的硬件配置选择最佳参数（如块大小、并行线程数等）。
    • 自动调优使得代码可以在不同的硬件平台上实现接近最优的性能。

triton.language里面能用的函数是一些基础款，也够用了，几乎没有学习成本

Triton的优势： 可以直接控制 SRAM(onchip, shared memory) 和 HBM(global memory) GPU memory

学着学着不对劲，还是先学LLVM吧。

Thanks for your support.