在過去的十年中，CUDA已經成為深度學習模型推理領域的核心技術。NVIDIA的CUDA技術優化了GPU的性能，提高了數據處理的速度和效率。然而，CUDA的主導地位也帶來了對NVIDIA硬件的依賴，這限制了在非NVIDIA平臺上進行高效推理的可能性。為了應對這一挑戰，近年來，開源社區和企業開始尋求開發不依賴于特定供應商的推理技術。

Triton的崛起

OpenAI開發的Triton編程語言應運而生，旨在為廣泛的硬件設備提供高效的并行編程方式。

它允許開發者編寫一次代碼，就能在NVIDIA、AMD甚至未來的Intel GPU上運行，極大地擴展了AI模型的適用范圍。Triton的出現，不僅打破了CUDA的壟斷，還為AI社區帶來了新的活力和創新的可能性。

與CUDA相比，Triton提供了一個更靈活、更高級的抽象層，使得開發者可以在多種類型的GPU上實現優化的計算性能，而不僅僅是NVIDIA的產品。

Triton的核心優勢在于其對于大型模型如LLM的支持。Triton能夠處理諸如矩陣乘法、正規化（RMSNorm）、自注意力機制（RoPE）、激活函數（如SiLU）以及元素級乘法等關鍵操作，這些都是LLM推理中不可或缺的部分。PyTorch社區已經利用Triton重寫了這些操作的內核，通過torch.compile實現了自動生成Triton內核的能力。

技術實現細節

將CUDA內核替換為Triton內核是一個復雜的過程，需要精細的調優和定制。例如，PyTorch團隊開發了一種自定義的FP16 Triton GEMM（General Matrix-Matrix Multiply）內核，該內核專門針對LLM中的線性層的矩陣乘法進行優化。這個過程涉及到對不同線性層權重矩陣形狀的精確調整，以確保最佳性能。

此外，flash attention是LLM中另一個關鍵操作，PyTorch團隊采用了多種配置的Triton flash attention內核進行實驗，最終選擇了AMD的Triton flash attention內核，因為它能夠在eager和compile模式下提供清晰的輸出，同時兼容torch.compile。

性能評估

在性能方面，雖然Triton內核的性能尚未完全達到CUDA的水平，但已經非常接近。在NVIDIA的H100和A100 GPU上的測試顯示，使用Triton內核的模型推理性能達到了CUDA模型性能的78%到82%。這一結果足以證明，即使是在高要求的應用場景下，非CUDA推理也能達到實用的水平。

未來發展方向

展望未來，Triton內核還有很大的優化空間。PyTorch團隊和OpenAI已經在探索更高效的內核設計，如利用更精細的工作分解策略和持久內核技術，這可能會進一步提高Triton的計算效率并縮小與CUDA的性能差距。此外，隨著FlexAttention和FlashAttention-3等新技術的引入，PyTorch有望在更廣泛的硬件平臺上實現更長的上下文長度和更大規模的解碼任務，這將推動整個行業向前發展。

Triton中國社區的發展和影響

2024年，智源研究院舉辦了多期Triton中國生態Meetup活動，這些活動不僅聚焦于技術分享和經驗交流，也致力于構建一個開放、協作的Triton技術生態體系。國內大量支持Triton語言編程的芯片廠商積極參與了交流。

隨著Triton和其他非CUDA技術的成熟，社區正逐步克服對特定硬件供應商的依賴，這將使得大模型推理技術更加普及和可用。這一進展不僅對科研人員和開發者具有重要意義，也對推動AI技術的普及具有深遠影響。隨著更多的創新和研究，未來的大模型推理將更加高效、靈活和開放。

參考鏈接：

1. https://pytorch.org/blog/cuda-free-inference-for-llms/
2. https://hub.baai.ac.cn/view/39486

本文轉載自公眾號AI時代窗口 作者：郁愈

原文鏈接：??https://mp.weixin.qq.com/s/TrzwSBQ301Grcpye9-pt1Q???