在過(guò)去的十年中，CUDA已經(jīng)成為深度學(xué)習(xí)模型推理領(lǐng)域的核心技術(shù)。NVIDIA的CUDA技術(shù)優(yōu)化了GPU的性能，提高了數(shù)據(jù)處理的速度和效率。然而，CUDA的主導(dǎo)地位也帶來(lái)了對(duì)NVIDIA硬件的依賴(lài)，這限制了在非NVIDIA平臺(tái)上進(jìn)行高效推理的可能性。為了應(yīng)對(duì)這一挑戰(zhàn)，近年來(lái)，開(kāi)源社區(qū)和企業(yè)開(kāi)始尋求開(kāi)發(fā)不依賴(lài)于特定供應(yīng)商的推理技術(shù)。

Triton的崛起

OpenAI開(kāi)發(fā)的Triton編程語(yǔ)言應(yīng)運(yùn)而生，旨在為廣泛的硬件設(shè)備提供高效的并行編程方式。

它允許開(kāi)發(fā)者編寫(xiě)一次代碼，就能在NVIDIA、AMD甚至未來(lái)的Intel GPU上運(yùn)行，極大地?cái)U(kuò)展了AI模型的適用范圍。Triton的出現(xiàn)，不僅打破了CUDA的壟斷，還為AI社區(qū)帶來(lái)了新的活力和創(chuàng)新的可能性。

與CUDA相比，Triton提供了一個(gè)更靈活、更高級(jí)的抽象層，使得開(kāi)發(fā)者可以在多種類(lèi)型的GPU上實(shí)現(xiàn)優(yōu)化的計(jì)算性能，而不僅僅是NVIDIA的產(chǎn)品。

Triton的核心優(yōu)勢(shì)在于其對(duì)于大型模型如LLM的支持。Triton能夠處理諸如矩陣乘法、正規(guī)化（RMSNorm）、自注意力機(jī)制（RoPE）、激活函數(shù)（如SiLU）以及元素級(jí)乘法等關(guān)鍵操作，這些都是LLM推理中不可或缺的部分。PyTorch社區(qū)已經(jīng)利用Triton重寫(xiě)了這些操作的內(nèi)核，通過(guò)torch.compile實(shí)現(xiàn)了自動(dòng)生成Triton內(nèi)核的能力。

技術(shù)實(shí)現(xiàn)細(xì)節(jié)

將CUDA內(nèi)核替換為T(mén)riton內(nèi)核是一個(gè)復(fù)雜的過(guò)程，需要精細(xì)的調(diào)優(yōu)和定制。例如，PyTorch團(tuán)隊(duì)開(kāi)發(fā)了一種自定義的FP16 Triton GEMM（General Matrix-Matrix Multiply）內(nèi)核，該內(nèi)核專(zhuān)門(mén)針對(duì)LLM中的線(xiàn)性層的矩陣乘法進(jìn)行優(yōu)化。這個(gè)過(guò)程涉及到對(duì)不同線(xiàn)性層權(quán)重矩陣形狀的精確調(diào)整，以確保最佳性能。

此外，flash attention是LLM中另一個(gè)關(guān)鍵操作，PyTorch團(tuán)隊(duì)采用了多種配置的Triton flash attention內(nèi)核進(jìn)行實(shí)驗(yàn)，最終選擇了AMD的Triton flash attention內(nèi)核，因?yàn)樗軌蛟趀ager和compile模式下提供清晰的輸出，同時(shí)兼容torch.compile。

性能評(píng)估

在性能方面，雖然Triton內(nèi)核的性能尚未完全達(dá)到CUDA的水平，但已經(jīng)非常接近。在NVIDIA的H100和A100 GPU上的測(cè)試顯示，使用Triton內(nèi)核的模型推理性能達(dá)到了CUDA模型性能的78%到82%。這一結(jié)果足以證明，即使是在高要求的應(yīng)用場(chǎng)景下，非CUDA推理也能達(dá)到實(shí)用的水平。

未來(lái)發(fā)展方向

展望未來(lái)，Triton內(nèi)核還有很大的優(yōu)化空間。PyTorch團(tuán)隊(duì)和OpenAI已經(jīng)在探索更高效的內(nèi)核設(shè)計(jì)，如利用更精細(xì)的工作分解策略和持久內(nèi)核技術(shù)，這可能會(huì)進(jìn)一步提高Triton的計(jì)算效率并縮小與CUDA的性能差距。此外，隨著FlexAttention和FlashAttention-3等新技術(shù)的引入，PyTorch有望在更廣泛的硬件平臺(tái)上實(shí)現(xiàn)更長(zhǎng)的上下文長(zhǎng)度和更大規(guī)模的解碼任務(wù)，這將推動(dòng)整個(gè)行業(yè)向前發(fā)展。

Triton中國(guó)社區(qū)的發(fā)展和影響

2024年，智源研究院舉辦了多期Triton中國(guó)生態(tài)Meetup活動(dòng)，這些活動(dòng)不僅聚焦于技術(shù)分享和經(jīng)驗(yàn)交流，也致力于構(gòu)建一個(gè)開(kāi)放、協(xié)作的Triton技術(shù)生態(tài)體系。國(guó)內(nèi)大量支持Triton語(yǔ)言編程的芯片廠(chǎng)商積極參與了交流。

隨著Triton和其他非CUDA技術(shù)的成熟，社區(qū)正逐步克服對(duì)特定硬件供應(yīng)商的依賴(lài)，這將使得大模型推理技術(shù)更加普及和可用。這一進(jìn)展不僅對(duì)科研人員和開(kāi)發(fā)者具有重要意義，也對(duì)推動(dòng)AI技術(shù)的普及具有深遠(yuǎn)影響。隨著更多的創(chuàng)新和研究，未來(lái)的大模型推理將更加高效、靈活和開(kāi)放。

參考鏈接：

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

本文轉(zhuǎn)載自公眾號(hào)AI時(shí)代窗口 作者：郁愈

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