九源开源技术周 04 | 九齿多后端,让简单的算子代码跑得更快

上一篇,我们把九齿矩阵乘里那一串难读的数据排布简化了。代码更好写了,接下来绕不开的问题是:跑得怎么样?

九齿一直希望大家按串行的思路写计算,把并行化交给编译器。但代码写得简单,后面的实现也得把硬件用好。过去,九齿主要通过 Triton 生成算子。在实际使用中,我们发现,一些算子,特别是在部分国产平台上,性能会受到 Triton 后端的限制。上层已经把计算描述清楚了,后端生成的代码却未必能充分发挥硬件能力。

为此,我们为九齿增加了 TileLang、CUDA C、Ascend C 和 BANG C 后端,让编译器有了更多生成实现的方式,针对不同硬件继续挖掘性能。

性能还取决于后端怎样生成代码

同样是矩阵乘,公式没有变,实际执行却有很多讲究:

  • 数据怎样成块搬运?
  • 搬数据和计算能不能重叠?
  • 块乘法能不能用好硬件的矩阵计算单元?

这些工作完成得怎么样,都会直接影响最后的性能。

如果所有实现都经过同一条编译路径,九齿能做出的选择也会受这条路径约束。当后端没有用好某些硬件能力时,单靠调整上层代码,未必就能补齐这部分差距。多后端让九齿可以利用不同编译体系的能力,也可以直接面向目标硬件生成原生代码。

计算只写一次,后端可以生成不同实现

在编程接口上,算子仍然由 排布 和 计算 两部分组成。拿最简单的向量加法来说,同一份定义可以分别交给 Triton 和 TileLang:

import ninetoothed as nt

def arrangement(x, y, out):
    return tuple(t.tile((256,)) for t in (x, y, out))

def application(x, y, out):
    out = x + y

tensors = tuple(nt.Tensor(1) for _ in range(3))

triton_add = nt.make(
    arrangement, application, tensors, backend="triton"
)

tilelang_add = nt.make(
    arrangement, application, tensors, backend="tilelang"
)

前面的排布和计算只定义一次,生成时指定后端。这样,我们可以保留同一份计算定义,在目标硬件上比较不同后端的表现,再确定采用哪种实现。

一套前端和中端,支撑多个后端

多后端能够共用同一份算子定义,背后是九齿对 前端、中端和后端 的清晰划分:

  • 前端: 负责读取排布和计算,整理成统一的中间表示;
  • 中端: 在这份表示上做公共分析和优化;
  • 后端: 结合目标硬件安排执行方式,生成对应语言的代码。

图 1 九齿编译器的职责分层

比如消除重复计算、复用归约结果,这些公共优化可以在中端实现,再被多个后端沿用。中端完善了一项优化,接入这条编译流程的后端就能复用相应成果,不必为每种输出语言各做一遍。

新增后端,也就不必从头实现一套编译器。前端对程序的理解、中端积累的公共优化都可以继续使用,后端开发集中在怎样把这些计算映射到目标硬件。

既能接上 TileLang,也能直接生成原生代码

1. TileLang 后端

以矩阵乘为例,九齿会分析输入怎样被访问、各个块怎样参与计算。在满足相应条件时,生成 T.copy 搬运数据块、T.Pipelined 组织流水循环、T.gemm 完成块矩阵乘,再由 TileLang 完成后续编译。

开发者写的是分块计算,数据搬运与计算流水等实现细节可以由编译器接手。[1]

📢 联合发布预告
在10月31日举办的 2026 GAI+C 大会的九源智能计算系统创新与开放生态论坛上,我们还与 TileLang 进行了联合发布,介绍了九齿与 TileLang 结合的最新进展。

2. 原生后端

原生后端则让九齿可以直接安排目标硬件上的执行方式。比如通过 Ascend C 和 BANG C,针对相应芯片组织数据搬运、向量计算和矩阵计算,让国产平台上的优化有更直接的实现空间。

把有效的优化留在编译器里

比如用户写一行求和,编译器可以让多个线程各算一部分,再把结果合起来。怎样分配这些工作、怎样组织中间结果,不必在每个算子里重新写一遍。九齿的多后端设计已经能为不同语言生成合适的归约表达实现。

我们希望九齿编译器能够让算子开发更加高效:

  • 大模型负责设计算法、写出简洁的九齿代码;
  • 九齿编译器负责生成高效算子;
  • 大模型在编译和运行过程中发现的优化方法,经验证后沉淀为九齿编译器的优化规则,供后续算子直接复用。

这样,大模型探索的优化能力就能不断积累为编译器的能力,而编译器也能减少大模型优化过程中的重复试错。

🔗 相关链接

【1】:
https://github.com/InfiniTensor/ninetoothed/blob/ab0c58c23f82ad94143066239f70a1d5ae941089/src/ninetoothed/backends/emitters/tilelang.py

项目仓库:
https://github.com/InfiniTensor/ninetoothed

关注与交流