Triton Kernel 优先：全新 LLM 推理方式

摘要：目前主流的推理框架如vLLM、TensorRT-LLM、SGLang都在MLsys上对大模型推理进行了优化，能够得到一个非常好的多并发成绩。triton_llm 项目带来了一个全新的llm推理方式，以 Triton Kernel为中心，希望能在batch较小时获得更好的性能。目前在 RTX 5090（Blackwell）上，针对 gpt-oss，已实现了与 Ollama 相当的单卡性能。

项目地址：OpenMLIR/triton_llm

一、项目背景与动机

这个项目在我跟着猛猿老师学了一段时间vLLM大模型推理框架并写了不少Triton kernel后就诞生了这个想法。

之后我开发了triton_runner，一个面向 Triton 的多层级 runner 工具，用于调试 Triton IR，支持在多个编译阶段直接运行 GPU kernel，包括 ttir、ttgir、llir、ptx、cubin，目的是想插入Triton的编译流程去进行性能调优和调试。

开发的过程中了解到Triton upstream的Gluon，一个复用Triton编译流程可以进行layout、调度、shared memory的语言，依然和Triton一样tile-based SPMD，我的triton_runner有了更大的用武之地。

8月6日，OpenAI再次回归开源，发布了gpt-oss，甚至提供了Triton实现，也就是可以使用Triton进行大模型的推理，这非常激动人心，我也在当天在修复了一个bug后，就在5090上跑通了，详见5090显卡+Triton，轻松玩转GPT-OSS-20B！，速度还不错。周末我顺手给Ampere架构也做了支持，手写一段内联ptx，在A800上测试了120B模型，在4090和3090上测试了20B模型，并给Triton提了PR7805。

8月13日，TritonDistributed-MegaTritonKernel发布，坚定了我做好我的项目的决心。

前段时间，triton_runner有了一个issue，让我开始思考我的工具给谁用的问题。先给我自己用！triton_llm就是我的答案，我可以使用triton_runner把gpt-oss推理性能提上来，并提供些使用triton_runner调试Triton kernel的实战教程来推广我的工具。有点为了”为了这碟醋包了这顿饺子“的感觉，但是实际不是的，程序员需要开发出有用的东西，有用才是开始。

之前在做triton_ocl(原triton-spirv)，面向OpenCL的Triton编译器后端，我把项目背景与动机写在了最后，现在我发现这才是开始，特别是开源项目。再贴一下那篇文章中的 Introduction / Motivation。

项目的源起在2023年5月，当时我在将我们自有的gpgpu芯片接入Pytorch，希望完成对大模型的训练，此时我了解到了Triton。在自有AI芯片接入AI框架Pytorch的方案中，我记录了我对Triton的看好。2024年6月随着智源研究院的推动，Triton也在国内火了起来。我也心痒痒，然后写了浅析 Triton 执行流程，并在文章中提到了想做一个SPIR-V后端的玩具。今年我正式投入到了Triton的开发中，空闲之余我也进行了算子开发的学习，并写了LeetGPU入门教程 (CUDA guide最佳实践)。我慢慢意识到了Triton的伟大，而不仅是初次了解时的扩充生态以及快速交差的期望。Pythonic is future!，世界上会写Python的人越来越多，AI模型的发展可能会产生更多的算子，让算法的人能直接写算子带来的收益很可能抹平算子性能的差距。我用业余时间维护triton-spirv一方面是为了完成当年想做一个toy的心愿，另外也希望降低Triton的接入成本，为大家提供一个简明教程。

triton_llm 将针对RTX 5090(Blackwell)在gpt-oss上进行深度优化。其他架构仅展示用，由于是Triton，可能有意外惊喜。

二、项目进展

其实是从8月6日gpt-oss发布时才开始开发的，花了几天熟悉了下gpt-oss中的infra，并写了如何评价gpt-oss的回答。

对sm120（RTX 5090、RTX PRO 6000 等）、sm90（H100、H200、H20）、sm80（A800、A100）、sm89（RTX 4090、RTX 6000、L40）、sm86（RTX 3090、A10 等）等多种架构进行了测试，Triton和Ollama有20 tokens/s的差距。

在做了profile后，这两天手搓了gpt-oss的RMSNorm 和 ROPE的Triton kernel，已经达到了和 Ollama相同的水平，在5090上你可以使用 triton_llm 得到111token/s的输出，可以运行triton_llm中的bench/bench_chat.py 得到Inter-Token Latency (ITL)。

目前面前还有一座大山，llama.cpp。OpenAI’s New Open Models Accelerated Locally on NVIDIA GeForce RTX and RTX PRO GPUs，256 token/s，so fast。

一些其他的数据对比，H20在20B模型在triton_llm是104token/s，在120B模型是70token/s。可以对比的是在H20使用vllm的单并发是150tokens/s，100并发3455tokens/s。并发越多vllm框架带来的惊喜越大，这也是MaaS(Model as a Service)平台提供低价优质服务的关键，打满打满。vllm数据来源于evalscope。

三、项目安装及使用

1、支持的 GPU

• sm120：RTX 5090、RTX PRO 6000 等
• sm90：H100、H200、H20 等
• sm80：A800、A100
• sm89：RTX 4090、RTX 6000、L40 等
• sm86：RTX 3090、A10 等

2、显存要求

• 若 GPU 显存 ≥ 24 GB，可运行 gpt-oss-20b。
• 若 GPU 显存 ≥ 80 GB，可运行 gpt-oss-120b。

3、安装

pip install torch==2.8.0
git clone https://github.com/OpenMLIR/triton_llm
cd triton_llm
pip install -e .[triton]
pip install -e triton_kernels

4、下载模型

modelscope很好用，速度也很快。下载120b模型请自行修改命令。

pip install modelscope

modelscope download openai-mirror/gpt-oss-20b  --include "original/*" --local_dir gpt-oss-20b/

5、运行

使用120b模型请自行修改命令。

# 测试
python -m gpt_oss.generate gpt-oss-20b/original/

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
# 对话
python -m gpt_oss.chat gpt-oss-20b/original/

6、性能测试

我目前在尝试优化 Inter-Token Latency(ITL)，即生成每个输出 token 的时间间隔，用来评估模型的生成速度。

python -m bench.bench_chat gpt-oss-20b/original/

# 展示输出，实验性质
python -m bench.only_output gpt-oss-20b/original/

四、项目展望

目前处于极早期，roadmap没有的。

计划从vLLM、TensorRT-LLM、SGLang里学习一些关键优化，加速现在推理速度。

另外也会参考TritonDistributed-MegaTritonKernel进行多卡的推理加速。

Pytorh的@torch.inference_mode()和NVIDIA的CUDA graph非常激动人心，提供了很好的性能。但是还是会有开销的且Pytorh依赖有点重，大概率会去掉。

利用KTransformers 的理念让显存不够的卡借助cpu也能跑起来 大概是很后期的事情了，短时间内不会启动。