2026 年 1 月 19 日，2025 冬季 InfiniTensor 大模型与人工智能系统训练营 Triton&九齿系列第二课《九齿一重：初窥》开讲。

本节课将详细讲解九齿语言的基本概念与设计初衷，以及九齿在深度学习神经网络领域的应用。

九齿概述

1. 九齿简介：

• 九齿是启元内部开发的一门提供张量级抽象的领域特定语言。
• 定位与 Triton 相似，用于编写深度学习神经网络领域的算子。
• 使用 Triton 作为编译后端，旨在简化编程并避免重复造轮子。

2. 九齿与 Triton、CUDA 的对比

• 九齿进一步抽象了并行编程细节，降低了并行编程门槛。
• Triton 相对于 CUDA 做了大量抽象，但仍需理解并行编程基础。
• CUDA 是底层语言，需要处理硬件细节和并行编程。

九齿语言实践

1. 向量加法示例

• 九齿与 Triton 对比

下图中使用 OpenAI Triton 语言实现向量加法（Vector Addition）的代码示例，这通常被视为 GPU 编程领域的“Hello World”。
与传统 CUDA C++ 开发模式不同，这段代码演示了如何利用 Python 语法配合 @triton.jit 装饰器来编写高性能内核：代码逻辑不再纠结于单个线程（Thread）的控制，而是以块（Block）为单位进行编程。它通过获取 program_id 来计算内存偏移量，利用 mask 掩码机制优雅地处理数组边界（防止内存越界），并使用 tl.load 和 tl.store 对数据块进行直接读写。这直观地体现了 Triton 的核心价值：通过更高层级的抽象，将开发者从复杂的线程束管理和共享内存同步中解放出来，同时依然保持极高的执行效率。

下图展示了使用 Ninetoothed（九齿） 框架实现向量加法的代码示例，这代表了一种比 Triton 抽象层级更高、更偏向声明式的编程范式。代码的核心亮点在于将 数据编排（Arrangement） 与 计算逻辑（Application） 进行了彻底解耦：开发者只需在 arrangement 函数中通过 .tile 声明数据的分块策略，在 application 函数中描述纯粹的数学运算（input + other），而无需像在 Triton 中那样手动处理指针偏移或边界掩码（Mask）。这种设计思路旨在通过编译器自动推导最优的并行实现，让开发者能够专注于算法定义而非底层硬件的内存管理。

• 张量可视化

• 排布可视化

• 九齿实现

下图展示了 Ninetoothed（九齿） 编程模型中算法逻辑与代码实现的直接映射关系。左侧列出了向量加法的算法步骤，其中红色方框圈出的部分描述了将向量A 、B 、C 划分为大小为 M 的数据块（Blocks）的逻辑；右侧代码中的红色方框则对应展示了如何在 arrangement 函数中通过 .tile((BLOCK_SIZE,)) 方法来实现这一分块策略。这种左右对照直观地说明了 Ninetoothed 如何让开发者将数学上的“数据分块”思想直接转化为代码中的“Tiling”声明，从而清晰地实现了数据编排（Arrangement）与计算逻辑（Application）的分离。

2. 矩阵乘法示例

• 算法描述：

下图直观地展示了高性能计算中核心的**分块矩阵乘法（Tiled Matrix Multiplication）**逻辑。图中将巨大的矩阵 （）和 （）划分成了微小的 BLOCK_SIZE 数据块，演示了如何通过计算 的行块与 的列块的乘积来生成结果矩阵 中的一个子块（绿色深色部分）。这是 GPU 编程（如 Triton 或 CUDA）优化 GEMM 算子的基础图景，其核心目的是通过分块提高数据局部性与复用率，使得计算过程中频繁访问的数据能够尽可能驻留在高速存储层级（如寄存器、片上 SRAM 或共享内存）中，从而减少对全局内存的访问。

分块矩阵乘法（Block Matrix Multiplication）是高性能计算中优化 GEMM（通用矩阵乘法）的核心逻辑。图中详细描述了如何将大规模的输入矩阵 A 和 B 切分为较小的子块（Blocks），并通过三层循环结构（ i,j确定输出块位置，k 进行累加归约）来计算结果矩阵 C 的每一个子块。这种算法设计的核心目的在于提高数据的局部性，使得计算过程中频繁使用的数据块能够被放入容量有限但速度极快的高速缓存（如 GPU 的 Shared Memory）中，从而避免频繁访问慢速的全局内存，突破内存带宽瓶颈以显著提升运算效率。

• 九齿实现：

Ninetoothed（九齿） 框架中实现矩阵乘法时至关重要的 “数据排布”（Arrangement） 代码通过定义 BLOCK_SIZE_M、N、K 来设定切块大小，并使用 .tile() 和 .expand() 操作符清晰地描述了数据在三维计算空间中的映射关系：它将输入矩阵（input）沿输出矩阵的列方向进行广播扩展，将权重矩阵（other）沿输出矩阵的行方向进行广播扩展。这种声明式的写法将复杂的分块索引计算与数据广播逻辑封装起来，替代了传统 CUDA 编程中繁琐的 blockIdx 和 threadIdx 映射计算，体现了该框架将“数据调度”与“数学计算”分离的核心设计理念。

Ninetoothed（九齿） 框架中实现矩阵乘法的最后一步：“应用与整合”（Application and Integration）。这里的 application 函数专注于描述纯粹的数学计算逻辑：它初始化一个累加器（Accumulator），通过循环遍历 维度（range(input.shape[0])）对分块后的输入矩阵和权重矩阵执行点积（ntl.dot）并累加，从而完成核心的矩阵乘法运算。图片底部展示了如何使用 ninetoothed.make 将定义好的数据排布（Arrangement）、计算逻辑（Application）以及张量形状描述整合在一起，生成最终的可执行 Kernel，这完美体现了该框架将“算法数学定义”与“底层数据调度”彻底解耦的设计哲学。

• 张量可视化

• 排布可视化

二维卷积

1. im2col：

使用 image to column 方法将卷积转换为矩阵乘。

• 优势：代码复用，减少编写量，提高开发效率。

2. implicit GEMM

在 Ninetoothed（九齿） 框架中实现 Implicit GEMM（隐式矩阵乘法），这是一种将卷积运算（Convolution）转化为矩阵乘法以利用高性能 GEMM 内核的经典优化技术。代码的核心在于自定义的 arrangement 函数：它通过一系列复杂的视图变换操作（如 .tile, .squeeze, .flatten, .permute），将高维的卷积输入（Input）和卷积核（Filter）在逻辑上“重塑”为矩阵乘法所需的二维形态。最精妙之处在于，最后一行 kernel = ... 直接复用了现有的矩阵乘法计算逻辑（mm.application），这生动地演示了该框架的强大之处——开发者只需定义数据的**“逻辑视图映射”**，就能直接利用现成的高性能计算单元，而无需手动编写复杂的 im2col 显存搬运代码。

3. FlashAttention

FlashAttention-2 的核心算法思想是通过分块（Tiling）和重计算（Recomputation）来减少对高带宽内存（HBM）的访问次数。图中展示了如何将 K（键）和 V（值）矩阵切分为不同的块（如$K^{(1)}, K^{(2)}$ ），使得庞大的中间结果——注意力分数S 和注意力权重A（图中橙色虚线框部分）——能够仅在 GPU 的片上高速存储层级中临时计算和存在，而无需写回慢速的 HBM。右侧的公式则展示了在线 Softmax（Online Softmax） 的更新机制，即通过维护统计量 来动态调整和重缩放（Rescaling）部分输出结果 ，从而在单次循环中精确计算出最终的注意力输出，极大地缓解了内存带宽瓶颈。

FlashAttention-2 前向传播（Forward Pass）的完整伪代码算法（Algorithm 1）从数学和逻辑层面揭示了其如何通过精细的显存管理来提升性能。算法的核心在于双层循环结构：外层循环遍历 Query 矩阵的分块（$mathbf{Q}_i$），将其加载到高速的片上内存（SRAM）中；内层循环则遍历 Key 和 Value 的分块（$mathbf{K}_j, mathbf{V}_j$），在 SRAM 内部直接计算注意力分数$mathbf{S}_i^{(j)}$ 并利用 Online Softmax 技术动态更新统计量（如最大值$m_i$ 和归一化因子 $ell_i$）以及输出结果$mathbf{O}_i$。这种设计确保了庞大的中间矩阵（如注意力图）从未离开过高速缓存，仅仅在最终计算完成后才将结果写入慢速的高带宽内存（HBM），从而大幅减少了 IO 访问瓶颈。

九齿相关信息

总结与展望

本节课讲了九齿语言的基本概念、核心特点和代码结构，以及九齿的相关应用。后续课程将会讲解更多九齿在深度学习领域的潜在应用，继续探索九齿语言的高级特性和优化技巧。

💡 学习贴士

更多细节及完整课程内容，请观看直播或查看课程回放。

链接：https://www.infinitensor.com/camp/winter2025

📚 训练营报名持续进行中

扫码报名与加入群聊，或点击 「阅读原文」 直达官网报名