Punica

問題定義 — k×n 個 GPU 的荒謬

LoRA 讓每個租戶都能用很低的成本微調出自己的模型。於是 ML 服務商面對的不再是「服務一個大模型」，而是「同時服務幾百個彼此不同、但共享同一個 backbone 的小模型」。用現有的 serving 系統做這件事，浪費得非常離譜。

如果把每個 LoRA 模型當成獨立訓練出來的模型來服務，就是在浪費 GPU。假設服務一個 LoRA 模型需要 k 個 GPU，服務 n 個不同的 LoRA 模型看起來就需要 k × n 個 GPU。這個直觀做法忽略了一件事：這些 LoRA 模型都源自同一個 pre-trained model，權重之間存在巨大的相關性。

三條設計準則

作者認為一個高效的多 LoRA 服務系統必須遵循三條準則，整篇論文的設計都是從這三條推出來的：

Punica 的三個貢獻

背景 — 為什麼是 decode 階段

要理解 Punica 為什麼長成這樣，得先理解 prefill 與 decode 這兩個階段在算術特性上的根本差異。這個差異是整篇論文的地基。

2.1 Prefill vs. Decode

Transformer LLM 的運作分兩階段。Prefill 接受使用者 prompt，生成後續的第一個 token 與一份 Key-Value cache（KvCache）。Decode 接受一個 token 與 KvCache，生成下一個 token 並在 KvCache 上追加一個 column。decode 是迭代過程，生成的 token 成為下一步的輸入，直到產生 end-of-sequence token。

假設 prompt 長度為 s、attention head 維度為 d：

差別一目了然：decode 階段的輸入是單一向量（batch 維度以外只有 1 個 token），GPU 使用率極低。

2.2 LoRA 的數學形式

LoRA 的關鍵觀察是：pre-trained model 與 fine-tune 後模型之間的權重差異具有低秩性。因此這個差異可以表示為兩個小而密的矩陣的乘積。

那麼，怎麼在共享 GPU 叢集上高效服務多租戶 LoRA？

一種做法是把每個 LoRA 模型當成獨立模型，用傳統 LLM serving 系統（例如 vLLM）服務。但這忽略了不同 LoRA 模型之間的權重共享——而這個共享正是大幅提升 GPU 效率的關鍵。更進一步，如果把每個 LoRA 模型視為獨立模型，模型載入時間會成為 bootstrap 階段的顯著效能瓶頸。

即使我們在各個 LoRA 模型之間共享 backbone，仍然留下一個問題：要怎麼高效地「批次計算」那些 LoRA add-on？這就是 SGMV 要解決的問題。

SGMV — 論文的技術核心

Segmented Gather Matrix-Vector Multiplication。名字很長，語意其實一句話就講完：把 batch 切成連續的 segment，每個 segment 乘上屬於它自己的那個權重矩陣。

3.1 問題的形式化

當一個 LoRA 模型在 batch 中有多個輸入時，可以把它們進一步 batch 在一起。作者把指向同一個 LoRA 模型的輸入連續地擺在一起。令 n 為 batch 中 LoRA 模型的數量，s_i 為第 i 個模型在 batch 中的最後一個元素索引，其中 s₀ = 0、s_n 為 batch size。輸入 {x_i | i ∈ [1, s_n]} 被切分為 {{x_j | j ∈ (s_{i-1}, s_i]} | i ∈ [1, n]}。

那麼 dense projection 的輸出可以寫成：

3.2 SGMV 的語意

注意 v += x·AB 這個 operator 可以拆成同一個 kernel 的兩次啟動：先初始化 v := 0，然後執行 v += x·A，接著 y += v·B。這個 operator 就叫 SGMV：

SGMV 把 batch 中不同請求的 feature-weight 乘法平行化，並把對應到同一個 LoRA 模型的請求 group 起來，藉此提高 kernel 的 operational intensity（算術強度），使其能夠使用 GPU Tensor Cores 加速。

3.3 Shrink 與 Expand：兩種截然不同的 kernel schedule

作者依據輸入與輸出的 feature 維度，把 SGMV operator 分成兩類：

兩個 kernel 都把 LoRA 索引綁定到 CUDA 的 blockIdx.y，於是每個 blockIdx.y 上的計算就是「features 與某個特定 LoRA 權重」之間的矩陣乘法。兩者都使用 GPU Tensor Cores 加速矩陣乘法。

3.4 Roofline 分析

作者用 roofline model 分析 SGMV kernel。FLOP 與記憶體 I/O bytes 的計算方式為：

• Distinct（每個請求都是不同 LoRA）：算術強度不會改變，因為 FLOP 與 I/O 以相同速率成長。由於每個輸入只用到少量 GPU 計算單元，加大 batch size 反而會提升效能（圖上是垂直往上）。
• Identical（全部同一個 LoRA）：曲線沿著記憶體頻寬的斜線往上爬，代表 SGMV 在這個情況下是被記憶體頻寬限制住的。
• Uniform / Skewed：介於兩者之間，是「平行度增加」與「算術強度增加」兩種效果的組合。

系統設計 — kernel 之外的四件事

有了 SGMV 只解決了「一張卡上怎麼算」。Punica 還需要回答「叢集上怎麼排」。這一節的四個機制——排程、on-demand 載入、請求遷移、KvCache layout——共同構成 Punica 的系統層。

4.1 排程新請求

Punica 的 scheduler 擁有所有 GPU 的全域狀態視圖。對每張 GPU，Punica 維護一個 working set（即 LLM invocation 的 batch 輸入）。隨著新請求加入 working set、decode step 展開，KvCache 消耗越來越多 GPU 記憶體，所以 Punica 也持續追蹤每張 GPU 上可用於 KvCache 的記憶體空間。

max batch size 上限在「叢集 throughput」與「per-token 延遲」之間取平衡。過大的 batch 會大幅拖慢延遲，卻只帶來邊際的 throughput 收益。作者 profile A100 之後，把 max batch size 定為 32。

這個 GPU 選擇邏輯刻意強調「在延遲約束內的叢集 throughput」，它的性質是：

• 忙碌的 GPU 傾向於維持忙碌，因為更多請求會被指派給它。
• 輕載的 GPU 傾向於降低負載，因為請求會逐漸終止。
• 閒置的 GPU 傾向於維持閒置。

因此 scheduler 能維持峰值 throughput，並依當前系統總負載整併 GPU 使用。這讓叢集的 scale up/down 更容易。在雲端環境中，Punica 執行兩種叢集配置操作：(1) 若叢集中沒有輕載 GPU，就請求更多 GPU；(2) 對於完全沒有負載的 GPU server，把資源歸還。

4.2 On-demand 模型載入

LoRA 模型與底層 pre-trained model 之間的權重共享讓模型載入非常快。LoRA 模型（矩陣 A 與 B）的大小只有底層 pre-trained model 的約 1%。

把 LoRA 模型從主記憶體載入 GPU 記憶體只是一次非同步的 host-to-device memory copy，延遲受 PCIe 頻寬限制：

4.3 請求遷移

隨著請求生成更多 token，它們的 KvCache 佔用越來越多 GPU 記憶體。當一張 GPU 的 KvCache 空間耗盡時，就把一些請求遷移到其他 GPU。請求遷移分兩步——evict 與 add。Scheduler 驅逐 GPU 上最新的請求，這保留了 FCFS 語意。被驅逐請求的排程方式與新增請求相同。

4.4 KvCache 的記憶體佈局

Punica 使用可分離（separable）的 KvCache layout，這對文字生成的 batching throughput 至關重要。HuggingFace Transformers 的 KvCache layout 由複雜的巢狀 tensor list 組成，概念上可視為：

這個佈局有兩個問題：

• 問題一（效率）：每個 decode step，HuggingFace Transformers 沿序列長度維度串接一個 tensor。這個串接很沒效率，因為它必須讀取整個 KvCache 並寫出一份新的副本，而新 tensor 只佔 1/S。
• 問題二（更嚴重，語意層面）：batching 維度不在最外層，這意味著 batch 中的請求在 KvCache 裡很難被分開。在這個限制下，一起進入 batch 的請求必須在所有 decode step 中維持在一起，直到所有請求都滿足各自的停止條件為止。

論文指出 FasterTransformer 與 DeepSpeed 也有類似問題。Punica 的 KvCache layout 則是：

作者使用 paged KvCache（PagedAttention）來最小化記憶體碎片化，並把 batching 維度放到最外層以支援 continuous batching。

實作與原始碼 — 論文沒說的那部分

Punica 的實作分兩部分：跑在單張 GPU 上的 PyTorch Python library，以及支援跨叢集服務的系統元件。這一節除了論文寫的內容，也直接引用 punica-ai/punica repo 裡的實際程式碼——有幾個地方跟論文的描述不完全一致，值得知道。

5.1 論文描述的實作

5.2 混合 Prefill 與 Decode 的批次策略

Punica 把處於 Prefill 階段的新請求與處於 Decode 階段的既有請求混在同一個 batch。這樣一來，dense projection 與 LoRA 都能受惠於更大的 batch size。

為了做 batching，所有輸入沿序列長度維度串接。Prefill 請求一律放在開頭，Decode 請求放在後半部。接著傳入一個 BatchLen struct 來區分不同請求：它包含每個 Prefill 請求的起始索引清單，以及 Decode 請求的數量。

在 self-attention 層，把索引與前導輸入狀態傳給 BatchPrefill kernel，把尾端的輸入狀態傳給 BatchDecode kernel。

為什麼是 7L？因為每一層有 7 個 LoRA add-on 需要計算。這一點在原始碼裡看得非常清楚。

5.3 原始碼：每一層的 7 次 LoRA 加法

下面是 src/punica/models/llama_lora.py 中 attention 層的 forward。注意 dense projection 照常呼叫，然後 add_lora() 就地（in-place）把 LoRA delta 加上去：

# src/punica/models/llama_lora.py
from punica.ops import add_lora_sgmv_custom_cutlass as add_lora

def forward(self, hidden_states, blen, prefill_kv, decode_kv,
            lora: BatchedLlamaLoraWeight):
    # 1) backbone 的 dense projection —— 一般的 batched GEMM
    q_proj = self.q_proj(hidden_states)
    k_proj = self.k_proj(hidden_states)
    v_proj = self.v_proj(hidden_states)

    # 2) LoRA add-on —— SGMV，就地累加到 q/k/v_proj 上
    add_lora(q_proj, hidden_states,
             lora.q.wa_ptr, lora.q.wb_ptr,
             lora.segment,          # ← 就是 SGMV 的 indptr s[]
             self.layer_idx, lora.rank)
    add_lora(k_proj, hidden_states,
             lora.k.wa_ptr, lora.k.wb_ptr, lora.segment,
             self.layer_idx, lora.rank)
    add_lora(v_proj, hidden_states,
             lora.v.wa_ptr, lora.v.wb_ptr, lora.segment,
             self.layer_idx, lora.rank)
    ...

Attention 有 q, k, v, o 四次，MLP 有 gate, up, down 三次 —— 合計每層 7 次，正好對應論文說的 7L。lora.segment 這一個 tensor 就是圖 2 裡的 s[]，整個 batch 的 LoRA 排程資訊全部濃縮在裡面。

5.4 原始碼：SGMV 的實際 API 與一個意外發現

repo 的 src/punica/ops/__init__.py 定義了對外的 kernel API。SGMV 的語意跟論文完全一致：

def sgmv_cutlass(y, x, w_ptr, s, layer_idx):
    """
    Semantics:
      y[s[i]:s[i+1]] += x[s[i]:s[i+1]] @ deref(w_ptr[i])

    Args:
      y: Shape: `[B, H2]`. Output vectors. Will be changed in-place.
      x: Shape: `[B, H1]`. Input vectors.
      w_ptr: Shape: `[S]`. DType: torch.int64. Pointer to the weight matrices.
      s: Shape: `[S+1]`, DType: torch.int32. Indptr of the weight matrices.
         `s[0] == 0`, `s[-1] == B`.
    """

但真正被模型使用的是這個 混合實作：

def add_lora_sgmv_custom_cutlass(y, x, wa_ptr, wb_ptr, s,
                                 layer_idx, lora_rank):
    """
    Semantics:
      y[s[i]:s[i+1]] += x[s[i]:s[i+1]] @ deref(wa_ptr[i]).T @ deref(wb_ptr[i])
    """
    tmp1 = torch.empty((8 * 1024 * 1024,), dtype=torch.uint8, device=x.device)
    tmp2_size = _kernels.sgmv_cutlass_tmp_size(wa_ptr.size(0))
    tmp2 = torch.empty((tmp2_size,), dtype=torch.uint8, device=x.device)
    v = torch.zeros((x.size(0), lora_rank), dtype=x.dtype, device=x.device)

    _kernels.sgmv_shrink(v, x, wa_ptr, s, tmp1, layer_idx)   # ← 自製 CUDA kernel
    _kernels.sgmv_cutlass(y, v, wb_ptr, s, tmp2, layer_idx)  # ← CUTLASS 實作

def sgmv(y, x, w_ptr, s, layer_idx):
    if x.size(1) < y.size(1):
        raise NotImplementedError("TODO: sgmv_expand")
    ...

權重的記憶體佈局也直接反映了這個混合設計 —— src/punica/utils/lora.py 的註解說得很白：

class LoraWeight:
    def __init__(self, num_layers, in_features, out_features,
                 lora_rank, dtype, device):
        # SGMV-Shrink custom CUDA kernel uses column-major.
        self.wa = torch.zeros(
            (num_layers, lora_rank, in_features), dtype=dtype, device=device)
        # SGMV-Expand cutlass kernel uses row-major.
        self.wb = torch.zeros(
            (num_layers, lora_rank, out_features), dtype=dtype, device=device)

5.5 原始碼：另一個 kernel — BGMV

repo 裡還有一個論文正文沒有命名、但確實對應到 §4 那段「每個請求都有不同 LoRA 索引」討論的 kernel：BGMV（Batched Gather Matrix-Vector），放在 csrc/bgmv/。

def add_lora_bgmv(y, x, wa_T_all, wb_T_all, indicies, layer_idx, scale):
    """
    Semantics:
      y[i] += (
          x[i].unsqueeze(0)
          @ wa_T_all[indices[i], layer_idx, :, :].transpose(-1, -2)
          @ wb_T_all[indices[i], layer_idx, :, :].transpose(-1, -2)
          * scale
        ).squeeze(0)

      indicies: Shape: `[B]`. Indices of the LoRA weights.
    """

實驗結果

6.1 實驗設定

四種請求分佈（決定 batch 中 LoRA 模型的多樣性）：

6.2 LoRA operator microbenchmark

作者把 batched LoRA operator 實作成兩次 SGMV kernel 啟動，並與兩個 PyTorch baseline 比較：

• Loop：對每個 LoRA 模型跑一次迴圈。
• Gather-BMM：gather 步驟把每個輸入需要的權重矩陣堆疊成單一矩陣，再用 torch.bmm() 做 batched 矩陣乘法。與 SGMV 一樣啟動 Gather 兩次、BMM 兩次。

作者的分析指出 Gather-BMM 的 I/O 比 SGMV 多得多：Gather 讀入 n × h_i × h_o 個元素、寫出 s_n × h_i × h_o 個；接著 BMM 又必須讀回 Gather 剛寫出的 s_n × h_i × h_o 個權重元素。合計 Gather-BMM 比 SGMV 多出 s_n × h_i × h_o × 2 個元素的記憶體 I/O。

6.3 不同 LoRA rank 的影響

這是論文中唯一一組數值最完整的實驗（Testbed #1，LoRA rank 8/16/32/64）：

6.4 Transformer layer benchmark

把 LoRA operator 併入 transformer layer 之後測整層效能（Testbed #1，7B 與 13B 配置，序列長度 512 與 2048）。

序列較短時 batching 效果更強：序列長度 512 時，batch size 從 1 增加到 32，延遲只增加 72%。序列較長時 self-attention 花的時間更多，削弱了 layer 層級的 batching 效果。

6.5 單 GPU 文字生成

Testbed #1 上生成 1000 個請求（約 101k tokens），所有系統以 FCFS 排程、max batch size 32。

論文的定性描述是：Punica 無論在哪種工作負載下都能穩定交出高 throughput。多數 baseline 在 Identical 情況下也能有相對高的 throughput，但一旦出現多個 LoRA 模型，效能就崩壞：

• Distinct：所有 baseline 系統的 batch size 都是 1，因此 throughput 很低。
• Uniform / Skewed：baseline 的多數 batch 的 batch size 都極小（1–3）。
• Punica：能把不同 LoRA 模型 batch 在一起，因此可以穩定跑在 batch size 32。

6.6 70B 模型 + Tensor Parallelism

在 Testbed #2 的 8 張 GPU 上，用 Megatron 的 tensor parallel scheme 跑 70B 模型。Punica 與 vLLM 比較，vLLM 也使用相同的 Megatron TP scheme。這組數字是論文中對比最鮮明的：

6.7 叢集部署

Testbed #2 的 16 張 GPU，7B 模型，實驗持續一小時。負載變化方式：宏觀來看請求率逐漸上升再逐漸下降；微觀來看，請求到達時間的間隔服從指數分佈，到達過程服從 Poisson 分佈。LoRA 模型熱門度服從 Zipf-1.5（與 Skewed 工作負載相同）。

論文的觀察是：

• Punica 能夠整併 GPU 使用率，同時交出高 throughput。
• GPU 通常在非閒置時以最大 batch size 運行，因為排程演算法優先考慮大 batch。
• 偶爾會有 GPU 以較小的 batch size 運行，因為它的 KvCache 空間耗盡而把一些請求遷移到其他 GPU。
• 當一張 GPU 變成閒置（batch size = 0），它很可能會維持閒置——必要時就能歸還給雲端供應商。

論文的 Figure 13 呈現三個面板（請求率、throughput、每張 GPU 的 batch size 隨時間變化），但未提供可讀取的原始數值，此處不重繪。

討論、限制與後續影響

7.1 作者明說的定位

7.2 從實驗設定看得出來的限制

7.3 這篇論文為什麼重要

Punica 是最早把「multi-tenant LoRA serving」當成獨立系統問題來處理的工作之一，而它的技術貢獻可以濃縮成一句話：把「不同模型不能一起 batch」這個看似根本的限制，用一個 kernel 拆掉了。

更關鍵的是 §6.4 那個實驗結論——LoRA add-on 的成本相對於 backbone 小到讓整層延遲對「batch 裡有幾種 LoRA」不敏感。這個性質一旦成立，多租戶 LoRA 服務的排程問題就從「模型放置問題」退化成「單純的 throughput 排程問題」，難度差了一個等級。Punica 系統層的所有設計，都是站在這個結論上的。

結論

Punica 的設計核心是一個新的 CUDA kernel（SGMV），它讓不同 LoRA 模型的 GPU 運算得以批次化。對每張 GPU，Punica 只需要底層 pre-trained model 的單一份副本，就能服務多個不同的 LoRA 模型，在記憶體與計算兩方面都顯著提升 GPU 效率。此外，Punica 的排程器在共享 GPU 叢集上整併多租戶 LoRA 服務工作負載。在固定大小的 GPU 叢集上，相較於 state-of-the-art LLM serving 系統，Punica 達到 12× 的 LoRA 模型服務 throughput，而每個 token 只增加 2ms 延遲。