一、TensorRT-LLM 的產品定位

TensorRT-LLM 是 NVIDIA 用于做 LLM（Large Language Model）的可擴展推理方案。該方案是基于 TensorRT 深度學習編譯框架來構建、編譯并執(zhí)行計算圖，并借鑒了許多 FastTransformer 中高效的 Kernels 實現(xiàn)，然后利用 NCCL 完成設備之間的通訊。考慮到技術的發(fā)展和需求的差異，開發(fā)者還可以定制算子來滿足定制需求，比如基于 cutlass 開發(fā)定制 GEMM。TensorRT-LLM 是一款致力于提供高性能并不斷完善其實用性的 NVIDIA 官方推理方案。

TensorRT-LLM 已經(jīng)在 GitHub 上開源，主要分為兩個分支，即 Release branch 和 Dev branch。其中 Release branch 每個月更新一次，而在 Dev branch 中則會較為頻繁地更新來自官方或社區(qū)中的功能，方便開發(fā)者體驗、評估最新功能。下圖展示了 TensorRT-LLM 的框架結構，其中除了綠色 TensorRT 編譯部分和一些涉及硬件信息的 kernels 外，其他部分都是開源的。

TensorRT-LLM 還提供了類似于 Pytorch 的 API 來降低開發(fā)者的學習成本，并提供了許多預定義好的模型供用戶使用。

考慮到大語言模型比較大，有可能單卡放不下，需要多卡甚至多機推理，因此 TensorRT-LLM 還提供了 Tensor Parallelism 和 Pipeline Parallelism 兩種并行機制來支持多卡或多機推理。

二、TensorRT-LLM 的重要特性

TensorRT-LLM 的重要特性之一就是豐富的模型支持。TensorRT-LLM 對主流大語言模型都提供了支持，比如 Qwen（千問）就是由開發(fā)者完成的模型適配，并已經(jīng)納入官方支持。用戶可以很容易地基于這些預定義的模型做擴展或定制。其二就是低精度推理，TensorRT-LLM 默認采用 FP16/BF16 的精度推理，并且可以利用業(yè)界的量化方法，使用硬件吞吐更高的低精度推理進一步推升推理性能。

另外一個特性就是 FMHA(fused multi-head attention) kernel 的實現(xiàn)。由于 Transformer 中最為耗時的部分是 self-attention 的計算，因此官方設計了 FMHA 來優(yōu)化 self-attention 的計算，并提供了累加器分別為 fp16 和 fp32 不同的版本。另外，除了速度上的提升外，對內存的占用也大大降低。我們還提供了基于 flash attention 的實現(xiàn)，可以將 sequence-length 擴展到任意長度。

如下為 FMHA 的詳細信息，其中 MQA 為 Multi Query Attention，GQA 為 Group Query Attention。

另外一個 Kernel 是 MMHA(Masked Multi-Head Attention)。FMHA 主要用在 context phase 階段的計算，而 MMHA 主要提供 generation phase 階段 attention 的加速，并提供了 Volta 和之后架構的支持。相比 FastTransformer 的實現(xiàn)，TensorRT-LLM 有進一步優(yōu)化，性能提升高達 2x。

另外一個重要特性是量化技術，以更低精度的方式實現(xiàn)推理加速。常用量化方式主要分為 PTQ(Post Training Quantization)和 QAT(Quantization-aware Training)，對于 TensorRT-LLM 而言，這兩種量化方式的推理邏輯是相同的。對于 LLM 量化技術，一個重要的特點是算法設計和工程實現(xiàn)的 co-design，即對應量化方法設計之初，就要考慮硬件的特性。否則，有可能達不到預期的推理速度提升。

TensorRT 中 PTQ 量化步驟一般分為如下幾步，首先對模型做量化，然后對權重和模型轉化成 TensorRT-LLM 的表示。對于一些定制化的操作，還需要用戶自己編寫 kernels。常用的 PTQ 量化方法包括 INT8 weight-only、SmoothQuant、GPTQ 和 AWQ，這些方法都是典型的 co-design 的方法。

INT8 weight-only 直接把權重量化到 INT8，但是激活值還是保持為 FP16。該方法的好處就是模型存儲2x減小，加載 weights 的存儲帶寬減半，達到了提升推理性能的目的。這種方式業(yè)界稱作 W8A16，即權重為 INT8，激活值為 FP16/BF16——以 INT8 精度存儲，以 FP16/BF16 格式計算。該方法直觀，不改變 weights，容易實現(xiàn)，具有較好的泛化性能。

第二個量化方法是 SmoothQuant，該方法是 NVIDIA 和社區(qū)聯(lián)合設計的。它觀察到權重通常服從高斯分布，容易量化，但是激活值存在離群點，量化比特位利用不高。

SmoothQuant 通過先對激活值做平滑操作即除以一個scale將對應分布進行壓縮，同時為了保證等價性，需要對權重乘以相同的 scale。之后，權重和激活都可以量化。對應的存儲和計算精度都可以是 INT8 或者 FP8，可以利用 INT8 或者 FP8 的 TensorCore 進行計算。在實現(xiàn)細節(jié)上，權重支持 Per-tensor 和 Per-channel 的量化，激活值支持 Per-tensor 和 Per-token 的量化。

第三個量化方法是 GPTQ，一種逐層量化的方法，通過最小化重構損失來實現(xiàn)。GPTQ 屬于 weight-only 的方式，計算采用 FP16 的數(shù)據(jù)格式。該方法用在量化大模型時，由于量化本身開銷就比較大，所以作者設計了一些 trick 來降低量化本身的開銷，比如 Lazy batch-updates 和以相同順序量化所有行的權重。GPTQ 還可以與其他方法結合使用如 grouping 策略。并且，針對不同的情況，TensorRT-LLM 提供了不同的實現(xiàn)優(yōu)化性能。具體地，對 batch size 較小的情況，用 cuda core 實現(xiàn)；相對地，batch size 較大時，采用 tensor core 實現(xiàn)。

第四種量化方式是 AWQ。該方法認為不是所有權重都是同等重要的，其中只有 0.1%-1% 的權重（salient weights）對模型精度貢獻更大，并且這些權重取決于激活值分布而不是權重分布。該方法的量化過程類似于 SmoothQuant，差異主要在于 scale 是基于激活值分布計算得到的。

除了量化方式之外，TensorRT-LLM 另外一個提升性能的方式是利用多機多卡推理。在一些場景中，大模型過大無法放在單個 GPU 上推理，或者可以放下但是影響了計算效率，都需要多卡或者多機進行推理。

TensorRT-LLM 目前提供了兩種并行策略，Tensor Parallelism 和 Pipeline Parallelism。TP 是垂直地分割模型然后將各個部分置于不同的設備上，這樣會引入設備之間頻繁的數(shù)據(jù)通訊，一般用于設備之間有高度互聯(lián)的場景，如 NVLINK。另一種分割方式是橫向切分，此時只有一個橫前面，對應通信方式是點對點的通信，適合于設備通信帶寬較弱的場景。

最后一個要強調的特性是 In-flight batching。Batching 是提高推理性能一個比較常用的做法，但在 LLM 推理場景中，一個 batch 中每個 sample/request 的輸出長度是無法預測的。如果按照靜態(tài)batching的方法，一個batch的時延取決于 sample/request 中輸出最長的那個。因此，雖然輸出較短的 sample/request 已經(jīng)結束，但是并未釋放計算資源，其時延與輸出最長的那個 sample/request 時延相同。In-flight batching 的做法是在已經(jīng)結束的 sample/request 處插入新的 sample/request。這樣，不但減少了單個 sample/request 的延時，避免了資源浪費問題，同時也提升了整個系統(tǒng)的吞吐。

三、TensorRT-LLM 的使用流程

TensorRT-LLM 與 TensorRT的 使用方法類似，首先需要獲得一個預訓練好的模型，然后利用 TensorRT-LLM 提供的 API 對模型計算圖進行改寫和重建，接著用 TensorRT 進行編譯優(yōu)化，然后保存為序列化的 engine 進行推理部署。

以 Llama 為例，首先安裝 TensorRT-LLM，然后下載預訓練模型，接著利用 TensorRT-LLM 對模型進行編譯，最后進行推理。

對于模型推理的調試，TensorRT-LLM 的調試方式與 TensorRT 一致。由于深度學習編譯器，即 TensorRT，提供的優(yōu)化之一是 layer 融合。因此，如果要輸出某層的結果，就需要將對應層標記為輸出層，以防止被編譯器優(yōu)化掉，然后與 baseline 進行對比分析。同時，每標記一個新的輸出層，都要重新編譯 TensorRT 的 engine。

對于自定義的層，TensorRT-LLM 提供了許多 Pytorch-like 算子幫助用戶實現(xiàn)功能而不必自己編寫 kernel。如樣例所示，利用 TensorRT-LLM 提供的 API 實現(xiàn)了 rms norm 的邏輯，TensorRT 會自動生成 GPU 上對應的執(zhí)行代碼。

如果用戶有更高的性能需求或者 TensorRT-LLM 并未提供實現(xiàn)相應功能的 building blocks，此時需要用戶自定義 kernel，并封裝為 plugin 供 TensorRT-LLM 使用。示例代碼是將 SmoothQuant 定制 GEMM 實現(xiàn)并封裝成 plugin 后，供 TensorRT-LLM 調用的示例代碼。

四、TensorRT-LLM 的推理性能

關于性能、配置等細節(jié)都可以在官網(wǎng)看到，在此不做詳細介紹。該產品從立項開始一直與國內很多大廠都有合作。通過反饋，一般情況下，TensorRT-LLM 從性能角度來說是當前最好的方案。由于技術迭代、優(yōu)化手段、系統(tǒng)優(yōu)化等眾多因素會影響性能，并且變化非常快，這里就不詳細展開介紹 TensorRT-LLM 的性能數(shù)據(jù)。大家如果有興趣，可以去官方了解細節(jié)，這些性能都是可復現(xiàn)的。

值得一提的是，TensorRT-LLM 跟自己之前的版本比，性能有持續(xù)地提升。如上圖所示，在 FP16 基礎上，采用了 KVQuant 后，速度一致的情況下降低了顯存的使用量。使用 INT8，可以看到明顯的吞吐的提升，同時顯存用量進一步降低?？梢?，隨著 TensorRT-LLM 優(yōu)化技術的持續(xù)演進，性能會有持續(xù)地提升。這個趨勢會持續(xù)保持。

五、TensorRT-LLM 的未來展望

LLM 是一個推理成本很高、成本敏感的場景。我們認為，為了實現(xiàn)下一個百倍的加速效果，需要算法和硬件的共同迭代，通過軟硬件之間 co-design 來達到這個目標。硬件提供更低精度的量化，而軟件角度則利用優(yōu)化量化、網(wǎng)絡剪枝等算法，來進一步提升性能。

TensorRT-LLM，將來 NVIDIA 會持續(xù)致力于提升 TensorRT-LLM 的性能。同時通過開源，收集反饋和意見，提高它的易用性。另外，圍繞易用性，會開發(fā)、開源更多應用工具，如 Model zone 或者量化工具等，完善與主流框架的兼容性，提供從訓練到推理和部署端到端的解決方案。

六、問答環(huán)節(jié)

Q1：是否每一次計算輸出都要反量化？做量化出現(xiàn)精度溢出怎么辦？

A1：目前 TensorRT-LLM 提供了兩類方法,即 FP8 和剛才提到的 INT4/INT8 量化方法。低精度如果 INT8 做 GEMM 時，累加器會采用高精度數(shù)據(jù)類型，如 fp16,甚至 fp32 以防止 overflow。關于反量化，以 fp8 量化為例，TensorRT-LLM 優(yōu)化計算圖時，可能動自動移動反量化結點，合并到其它的操作中達到優(yōu)化目的。但對于前面介紹的 GPTQ 和 QAT，目前是通過硬編碼寫在 kernel 中，沒有統(tǒng)一量化或反量化節(jié)點的處理。

Q2：目前是針對具體模型專門做反量化嗎？

A2：目前的量化的確是這樣，針對不同的模型做支持。我們有計劃做一個更干凈的api或者通過配置項的方式來統(tǒng)一支持模型的量化。

Q3：針對最佳實踐，是直接使用 TensorRT-LLM 還是與 Triton Inference Server 結合在一起使用？如果結合使用是否會有特性上的缺失？

A3：因為一些功能未開源，如果是自己的 serving 需要做適配工作，如果是 triton 則是一套完整的方案。

Q4：對于量化校準有幾種量化方法，加速比如何？這幾種量化方案效果損失有幾個點？In-flight branching 中每個 example 的輸出長度是不知道的,如何做動態(tài)的 batching？

A4：關于量化性能可以私下聊，關于效果，我們只做了基本的驗證，確保實現(xiàn)的 kernel 沒問題,并不能保證所有量化算法在實際業(yè)務中的結果，因為還有些無法控制的因素，比如量化用到的數(shù)據(jù)集及影響。關于 in-flight batching，是指在 runtime 的時候去檢測、判斷某個 sample/request 的輸出是否結束。如果是，再將其它到達的 requests 插進來，TensorRT-LLM 不會也不能預告預測輸出的長度。

Q5：In-flight branching 的 C++ 接口和 python 接口是否會保持一致？TensorRT-LLM 安裝成本高，今后是否有改進計劃？TensorRT-LLM 會和 VLLM 發(fā)展角度有不同嗎？

A5：我們會盡量提供 c++ runtime 和 python runtime 一致的接口，已經(jīng)在規(guī)劃當中。之前團隊的重點在提升性能、完善功能上，以后在易用性方面也會不斷改善。這里不好直接跟 vllm 的比較，但是 NVIDIA 會持續(xù)加大在 TensorRT-LLM 開發(fā)、社區(qū)和客戶支持的投入，為業(yè)界提供最好的 LLM 推理方案。