TritonLLM v0.1.0: 基于 Triton 的大模型推理初步探索

摘要：TritonLLM v0.1.0版本已经发布，已经可以pip install安装。gpt-oss-20b在5090上的decode速度从官方发布的90tokens/s来到了136tokens/s。本文记录了影响性能的关键优化方法，同时也总结了我在实践过程中一些未能奏效甚至错误的优化尝试。

项目地址：OpenMLIR/TritonLLM，另有短域名tritonllm.top方便访问。

一、项目进展

距离Triton Kernel 优先：全新 LLM 推理方式 已有半个月时间，最近做自动从ModelScope下载模型的支持，还做了RoPE、FP4以及部分heuristic优化，吐出速度从111tokens/s来到了136tokens/s。ollama在这些天也进行了更新，吐出速度来到了140token/s，当然离llama.cpp的256token/s都还很遥远。

二、项目使用

项目以RTX 5090(Blackwell)在gpt-oss的优化为中心，也支持4090、3090、A100等硬件，具体见项目README.zh.md。

1、安装

你可以通过 pip 安装 tritonllm 的最新稳定版本

pip install tritonllm

2、命令行界面 (CLI)

快速启动 gpt-oss-20b 模型的对话，将自动从 ModelScope 魔搭下载。

tritonllm

你也可以运行 tritonllm 120b 来直接使用 120B 模型。更多命令行参数请参考README.zh.md

三、关键优化技术

1、CUDA graph

2、MOE的FP4支持

3、RoPE和RMSNorm的手写算子

4、RoPE算子的block优化

4、triton_kernels的heuristic优化

4、sink Attn的选择优化

四、失败优化

1、版本升级

无论是triton_kernel的升级，还是Triton的升级目前都会给项目带来负收益。Triton upstream主要针对B200在做，5090的Blackwell还是弱了不少。

2、Liner替换

Torch的Liner算子非常厉害，特别是单batch下，Triton没有胜算，只能之后搞搞量化。

3、TMA attn

TMA在5090 单batch并不会更快，而且在CUDA graph下有些不稳定。

五、已知问题

1、IndexKernel Assert failed

/pytorch/aten/src/ATen/native/cuda/IndexKernel.cu:195: operator(): block: [0,0,0], thread: [96,0,0] Assertion idx >= 0 && idx < self_dim_size && "index_copy_(): index out of bounds" failed. 输出后整个chat就挂了，不知道是哪里的bug。偶发

2、o200k_base.tiktoken下载问题

有时候会因为网络问题，导致tiktoken下载失败。可以使用https://tritonllm.top/down/o200k_base.tiktoken下载并重命名为fb374d419588a4632f3f557e76b4b70aebbca790，然后设置下TIKTOKEN_CACHE_DIR即可。也可以取消utils.py的以下两行注释让其自动完成，其也可能存在网络问题，中断掉重开就行了。

   os.environ["TIKTOKEN_CACHE_DIR"] = tritonllm_bin_dir

    # download o200k_base.tiktoken
    save_file_to_tritonllm_bin_dir(tritonllm_bin_dir)