议程：

1. 通过线程束特化改进指令级并行 (ILP)

2. Triton-shared (进展和更新)

3. 关于通用张量描述符的问题

会议记录：

通过线程束特化改进指令级并行 (ILP)

主讲人：Hongtao Yu (Meta), Yuanwei (Kevin) Fang (Meta), Manman Ren (Meta)

备注

• Pytorch 2.6 及 Triton 发布分支 3.2

• 目标：Nvidia Hopper 架构，Blackwell 即将推出。

• 性能

  • Meta 的 FP8Rowwise GEMM (3-5% 提升，1D 持久循环)

  • FlashAttention (10-15% 提升，通过流水线和乒乓调度可能更快)。

• 什么是线程束特化？

  • 改进硬件指令调度。GPU 不具备良好的动态指令调度能力。

  • 使用多路线程束调度器。允许单个核心上的线程束针对不同的功能单元（例如内存、ALU、张量核心等）。所有这些都并行运行。

• 使用 GEMM 进行比较 * *

  • 统一线程束：8 个线程束，每个加载/处理 1/8 的数据。分成两组，每组处理 1/2 的数据。适用于 GEMM，但不适用于更复杂的内核。

  • 线程束特化：12 个线程束，其中 4 个线程束仅用于数据生成（只进行加载），8 个线程束仅用于 wgmma（只进行 wmma 操作）。为 Flash Attention 等更复杂的内核释放更多容量。

• 编译器实现

  • 如何启用线程束特化

    • 通过向自动调优配置添加两个开关来自动启用。

      • Num_consumer_groups - 非加载线程束组

      • Num_buffer_warp_spec - 生产者和消费者之间的缓冲区数量

  • 概念

    • 异步任务与其他异步任务并行运行。

    • 任务应使用不同的内存和 GPU 资源。

    • 通过共享内存和屏障进行同步协调。

  • 编译器实现

    • 自动任务分区。

    • 数据流多缓冲

  • 任务分区

    • 自动任务分区识别加载、ALU 操作、存储等任务。

    • 识别依赖链。将生产者链接到消费者。

    • 继续分区并在生产者和消费者线程束中插入同步原语。

  • 多缓冲

    • 生产者以循环方式继续加载/填充缓冲区，而消费者处理单个缓冲区。

    • 当没有可用空闲缓冲区时，生产者阻塞。

  • 未来展望

    • 多维循环的多缓冲

    • 单组中多个区域的缓冲区复用

    • 复杂的控制流，分区方案（乒乓调度，支持 Blackwell）

• 案例研究：Flash Attention - Kevin 和 Manman

  • 未使用 WS 时

    • 计算吞吐量：45%

    • 内存吞吐量：35%

    • SM 繁忙度：46%

    • 无交错：张量核心运行时 CUDA 核心空闲

  • 使用 WS 时

    • 计算吞吐量：69%

    • 内存吞吐量：35%

    • SM 繁忙度：71%

    • 交错 (加速原因)

      • TMA 与 CUDA 核心操作重叠

      • CUDA 核心与张量核心重叠

      • 张量核心与指令发布重叠。

    • 数据分区

    • 通信流水线和乒乓调度

    • 乒乓调度是一个命名屏障对。区域内只能有一个消费者。

问题

• 问：AMD 是否有等效的线程束组？这是否适用于 AMD GPU？

• 答：Meta 正在为 AMD 做这项工作。AMD 中没有命名屏障。正在通过在 AMD 上使用共享内存原子操作来模拟以获得相同的效果。

• 问：对于编译器难以检测的复杂情况，将这些机制提升到 Triton 内部的更高层级是否有意义？

• 答：是的。我们允许用户在 facebookexperimental/triton 中使用其分区来标注程序。我们希望看看是否能实现更多的自动化。

• 问：我们应该首先关注什么？线程束特化还是软件流水线作为初始优化？根据您的经验，哪种降级更可取？您打算将其带到主分支吗？

• 答：它们并非互斥。您需要找出对您而言有意义的部分。WS 优点：对流水线的外部循环支持。WS 优点：CUDA 核心与张量核心的重叠。

• 问：您看到了哪些改进？

• 答：Flash Attention：20% + 计算流水线和乒乓调度方法接近 Flash Attention v3 的性能。

Triton-shared (进展和更新)

演示者：Nhat Nguyen (Microsoft), Haishan Zhu (Meta)

备注

目标：

• 将 Triton IR 降低到 MLIR 核心方言 (linalg, memref, …)，更容易在 CPU 上运行。

• 专注于支持加速器的跨步内存访问

• 已开源于 https://github.com/microsoft/triton-shared

  • 努力使其与 OSS triton 保持同步（尽管略有延迟）

进展

• 编译器通道模块化。将数据提取与降级解耦。允许自定义降级流程。分析失败时行为可预测。

  • Triton-to-structured

  • triton-arith-to-linalg

  • Structured-to-memref

• 指针分析的改进

  • 支持嵌套循环

  • 非连续内存访问。

• 支持使用单个基指针的非结构化访问降级

• 支持将 Triton 操作降级到 linalg/mlir (如 split, join, cat 等)

路线图

• 完全支持非连续指针

• 检测其他内存访问模式（例如行收集/分散指针序列）

• 扩展到控制流操作

感谢！

Meta, 高通和社区

问题

• 问：未来的计划中，您想优先处理哪些项目？

• 答：许多 Triton 内核具有无法检测到的内存访问模式。我们没有回退解决方案（例如收集-分散支持）。需要等待 MLIR 指针方言落地才能使用。如果加载是连续的，MxN 加载指针分析会失败。但行可能是连续的，因此我们可以将分析分成多个块（行分散，行收集）。

• 答：在指针分析无法提取信息的地方，我们保持 IR 不变，以便现有的通道可以处理它们。我们可以处理指针张量上的循环迭代（常见模式）。像 if/else 这样更复杂的操作看起来是唾手可得的成果。

关于通用张量描述符的问题

• 问：通用张量描述符编程的进展如何？它不是 Nvidia 特有的。（上个月的问题）。

• 答：TMA 加速器可能会在各种 GPU 上变得更加通用。

• 答：TMA（张量描述符）支持预计将在未来几周内推出。将为不带 TMA 的 GPU 添加兼容模式（但可能会更慢）。并将添加块指针支持。我们将弃用主机端张量描述符（它们只为持久化内核提供了微小的性能提升）。允许用户自动调优。

会议纪要：

录音链接 此处