使用 Tiny-vLLM 從零開始構建高性能 LLM 推理引擎

現代大型語言模型 (LLMs) 的複雜性往往掩蓋了實際運行它們所需的基礎工程技術。雖然像 PyTorch 這樣的高階函式庫讓模型設計變得容易上手，但從權重文件轉換到響應迅速的推理伺服器，需要深入研究記憶體管理、GPU kernel 工程以及線性代數。

Tiny-vLLM 是一個旨在揭開這一過程神秘面紗的教學專案。它並非提供一個黑盒，而是作為一個引導式課程，教你如何使用 C++ 和 CUDA 在 Llama 3.2 1B Instruct 模型作為參考的情況下，構建一個高性能的推理引擎。它彌合了理論架構與硬體加速現實之間的差距。

推理伺服器的解剖結構

要構建一個推理引擎，首先必須理解 LLM 在物理層面上是什麼：一個包含數百萬個浮點數（權重）的文件，以及一個定義這些數字如何在一系列操作中被使用的藍圖（架構）。

推理伺服器的主要工作是將該藍圖轉換為可執行代碼。為了實現高性能，這通常使用 C++ 和 CUDA 來最大化硬體利用率。核心目標是在最小化延遲的同時處理多個提示詞 (prompts)，這需要將計算從 CPU 移至 GPU，在 GPU 上，矩陣乘法——LLM 的核心基礎——可以在數千個核心上進行並行化處理。

模型加載：Safetensors 與 BF16

大多數現代模型都以 Safetensors 格式分發。一個 Safetensors 文件由標頭大小、包含張量元數據（dtype、shape 和 offsets）的 JSON 標頭以及原始張量數據組成。

Bfloat16 的選擇

LLM 推理中的一個關鍵技術決策是數值格式。雖然標準的 16 位元浮點數 (FP16) 很常見，但許多模型使用 Bfloat16 (BF16)。

BF16 以精度換取範圍。它使用 8 位元用於指數部分（與 32 位元浮點數相同），而分數部分僅使用 7 位元。這使得它在 Transformer 層所需的巨量加總過程中，顯著降低了溢位 (overflow) 和欠位 (underflow) 的風險，且經驗證據表明，精度損失對於其提供的穩定性而言是可接受的權衡。

GPU 記憶體管理

高效的推理需要一套嚴格的數據移動策略，用於在 Host (CPU/DRAM) 與 Device (GPU/VRAM) 之間進行數據傳輸。由於 GPU 無法直接訪問系統 DRAM，數據必須透過 cudaMalloc 進行顯式分配，並透過 cudaMemcpy 進行傳輸。

為了優化性能，引擎遵循以下幾條黃金法則：

  1. 最小化分配： 預先分配大型緩衝區。
  2. 重複使用記憶體： 使用生命週期分析來識別在數據不再需要時可以重新利用的緩衝區。
  3. 減少傳輸： 盡可能減少將數據複製到 GPU 的次數。

工程化 CUDA Kernels

實現 Transformer 的前向傳播 (forward pass) 需要編寫自定義的 CUDA kernels。這涉及以 SIMT (Single Instruction, Multiple Threads) 的概念進行思考，其中相同的函數會在被分組為 blocks 和 warps 的數千個執行緒中執行。

Embedding Gather

第一步是將輸入的 token (整數) 映射到其對應的 embedding 向量。由於大多數 GPU 的每個 block 有 1024 個執行緒的限制，但 Llama 3.2 的 embedding 長度為 2048，因此一種常見的優化方法是讓每個執行緒處理 embedding 的兩個元素，從而有效地將每個 block 的吞吐量翻倍。

RMSNorm 與並行歸約 (Parallel Reduction)

均方根層歸一化 (RMSNorm) 需要計算整個 embedding 的平方平均值的平方根。這會產生同步挑戰：如何跨不同執行緒對數值進行加總？

Tiny-vLLM 使用 Parallel Reduction (Tree Reduction) 來實現這一點。執行緒不使用單一共享變數（這會導致競態條件），而是將其局部加總值寫入 __shared__ 記憶體向量中。然後，它們以樹狀結構迭代地加總這些值，每一步都將活動執行緒的數量減半，直到在索引 0 處剩餘一個總和。

為了確保數值穩定性並防止除以零（這會導致 NaN 並導致推理崩潰），在 RMS 計算中加入了一個 epsilon 值（例如 1e-05）。

矩陣乘法挑戰

矩陣乘法是引擎中計算量最大的部分。雖然可以編寫自定義 kernel，但業界標準是 cuBLAS，NVIDIA 的高度優化線性代數函式庫。

然而，存在一個顯著的摩擦點：cuBLAS 預期矩陣為 column-major format（列主序），而 LLM 權重通常以 row-major format（列主序）存儲。Tiny-vLLM 採用了一種數學轉置技巧來避免在記憶體中物理性地重新排列數據。透過操作轉置標誌 (CUBLAS_OP_T 和 CUBLAS_OP_N) 以及 GEMM (General Matrix Multiply) 調用中的操作順序，引擎可以將行主序數據視為列主序，從而顯著降低開銷。

優化吞吐量：批處理與緩存

Prefill 與 Decode

推理過程分為兩個截然不同的階段：

  • Prefill： 引擎處理整個輸入提示詞以生成第一個 token。這在計算上很重，但具有高度並行性。
  • Decode： 引擎逐一生成後續的 token。這這是一個受記憶體頻寬限制的過程，因為它只處理最後一個生成的 token。

KV Cache

為了避免在 decode 階段為每個先前的 token 重新計算 Key (K) 和 Value (V) 的投影，引擎實現了 KV Cache。此緩存存儲了所有已處理 token 的 K 和 V 向量，使得模型可以簡單地附加新 token 的投影，並對現有的歷史記錄進行注意力機制計算。

進階批處理 (Advanced Batching)

為了擴展伺服器規模，Tiny-vLLM 探索了兩種批處理方式：

  • Static Batching： 同時處理多個請求。雖然這會增加吞吐量，但會引入延遲，因為整個批次 (batch) 必須等待最長的序列完成。
  • Continuous Batching： 一種更先進的方法，當有新的插槽 (slot) 開放時，新請求可以立即進入批次，從而消除與靜態填充 (static padding) 相關的浪費。

結論

從零開始構建一個推理引擎揭示了，LLM 的「智能」並非源於單行代碼，而是源於參數的巨大規模以及計算的效率。透過實現這些底層層級的原始操作——從 RMSNorm 中的樹狀歸約到 cuBLAS 的轉置技巧——開發者可以超越僅將 LLM 作為一種服務來使用，並開始理解支撐現代 AI 的硬體與軟體協同設計。

Sources