在NVIDIA TensorRT-LLM中啟用ReDrafter的一些變化

Recurrent Drafting (簡稱 ReDrafter) 是蘋果公司為大語言模型 (LLM) 推理開發(fā)并開源的一種新型推測解碼技術(shù)，該技術(shù)現(xiàn)在可與 NVIDIA TensorRT-LLM 一起使用。ReDrafter 幫助開發(fā)者大幅提升了 NVIDIA GPU 上的 LLM 工作負(fù)載性能。NVIDIA TensorRT-LLM 是一個(gè) LLM 推理優(yōu)化庫，提供了一個(gè)易于使用的 Python API 來定義 LLM 和構(gòu)建 NVIDIA TensorRT 引擎，這些引擎具有頂尖的優(yōu)化功能，可在 GPU 上高效執(zhí)行推理。優(yōu)化功能包括自定義 Attention Kernel、Inflight Batching、Paged KV Caching、量化技術(shù) (FP8、INT4 AWQ、INT8 SmoothQuant) 等。

推測解碼 (Speculative decoding)是一種通過并行生成多個(gè) token 來加速 LLM 推理的技術(shù)。它使用較小的“draft”模塊預(yù)測未來的 token，然后由主模型進(jìn)行驗(yàn)證。該方法通過更好地利用可用資源實(shí)現(xiàn)低延遲推理，在保持輸出質(zhì)量的同時(shí)大大縮短了響應(yīng)時(shí)間，尤其是在低流量時(shí)段。

ReDrafter 運(yùn)用基于循環(huán)神經(jīng)網(wǎng)絡(luò) (RNN)的采樣 (稱為 Drafting) 并結(jié)合之前在 Medusa 等其他技術(shù)中使用的樹狀注意力，預(yù)測和驗(yàn)證來自多個(gè)可能路徑的draft token 以提高準(zhǔn)確性，并在解碼器的每次迭代中接受一個(gè)以上 token。NVIDIA 與蘋果公司合作，在 TensorRT-LLM 中添加了對該技術(shù)的支持，使更加廣泛的開發(fā)者社區(qū)能夠使用該技術(shù)。

ReDrafter 與 TensorRT-LLM 的集成擴(kuò)大了該技術(shù)的覆蓋范圍，解鎖了新的優(yōu)化潛力，并改進(jìn)了 Medusa 等先前的方法。Medusa 的路徑接受和 token 采樣發(fā)生在 TensorRT-LLM 運(yùn)行時(shí)，需要在接受路徑未知的情況下處理所有可能的未來路徑，而且其中大部分路徑最終都會(huì)被丟棄，這就給引擎內(nèi)部帶來了一些開銷。為了減少這種開銷，ReDrafter 要求在 drafting下一次迭代的未來 token 之前，先驗(yàn)證 token 并接受最佳路徑。

為了進(jìn)一步減少開銷，TensorRT-LLM 更新后在單個(gè)引擎中整合了drafting和驗(yàn)證邏輯，不再依賴運(yùn)行時(shí)或單獨(dú)的引擎。這種方法為 TensorRT-LLM 內(nèi)核選擇和調(diào)度提供了更大的自由度，通過優(yōu)化網(wǎng)絡(luò)實(shí)現(xiàn)了性能的最大化。

為了更好地說明 ReDrafter 的改進(jìn)，圖 1 展示了 TensorRT-LLM 中 ReDrafter 實(shí)現(xiàn)與 Medusa 實(shí)現(xiàn)的主要區(qū)別。大多數(shù)與推測解碼相關(guān)的組件都在 ReDrafter 的引擎內(nèi)完成，這大大簡化了 ReDrafter 所需的運(yùn)行時(shí)更改。

圖 1. NVIDIA TensorRT-LLM 中

Medusa（左）和 ReDrafter（右）實(shí)現(xiàn)的比較

下面將深入探討有助于在 TensorRT-LLM 中啟用 ReDrafter 的一些變化。

兼容 Inflight-batching

批處理的引擎

Inflight-batching (IFB) 是一種通過批量處理上下文階段和生成階段請求，來顯著提高吞吐量的策略。鑒于上下文階段請求與生成階段請求的處理方式不同（生成階段請求需要 draft token 驗(yàn)證），因此結(jié)合 IFB 的推測解碼會(huì)給管線帶來更大的復(fù)雜性。ReDrafter 將驗(yàn)證邏輯移至模型定義內(nèi)部，因此引擎在驗(yàn)證過程中也需要該邏輯。與注意力插件類似，該批處理被分成兩個(gè)較小的批處理：一個(gè)用于上下文請求，另一個(gè)用于生成請求。然后，每個(gè)較小的批處理進(jìn)入計(jì)算工作流，最后再合并成一個(gè)批處理進(jìn)行 drafting 流程。

圖 2. ReDrafter 兼容 TensorRT-LLM 引擎的

Inflight-batching 批處理計(jì)算工作流

請注意，這種方法要求任一路徑上的所有運(yùn)算符都支持空張量。如果一個(gè)批處理由所有上下文請求或所有生成請求組成，就可能出現(xiàn)空張量。該功能增加了 TensorRT-LLM API 的靈活性，使未來定義更復(fù)雜的模型成為可能。

實(shí)現(xiàn)引擎內(nèi)驗(yàn)證和 Drafting

為了在引擎內(nèi)進(jìn)行驗(yàn)證和 draft，TensorRT-LLM 更新時(shí)加入了對許多新操作的支持，這樣 PyTorch 代碼就可以輕松地轉(zhuǎn)化成一個(gè) TensorRT-LLM 模型的定義。

以下 PyTorch 代碼摘錄是蘋果公司的 PyTorch 實(shí)現(xiàn)的 ReDrafter。TensorRT-LLM 實(shí)現(xiàn)幾乎就是 PyTorch 版本的直接逐行映射。

PyTorch

def unpack(
  packed_tensor: torch.Tensor,
  unpacker: torch.Tensor,
) -> torch.Tensor:
  assert len(packed_tensor.shape) == 3
  last_dim_size = packed_tensor.shape[2]
  batch_size, beam_width, beam_length = unpacker.shape
  unpacked_data_indices = unpacker.view(
    batch_size, beam_width * beam_length, 1).expand(
    -1, -1, last_dim_size
  )
  unpacked_tensor = torch.gather(
    packed_tensor, 1, unpacked_data_indices).reshape(
    batch_size, beam_width, beam_length, -1
  )
  return unpacked_tensor

TensorRT-LLM

def _unpack_beams(
  x: Tensor,
  indices: Tensor,
  num_beams: int,
  beam_length: int
) -> Tensor:
  assert x.rank() == 3
  d0 = shape(x, 0, INT_DTYPE_STR)
  dl = shape(x, -1, INT_DTYPE_STR)
  indices = view(
    indices, [-1, num_beams * beam_length, 1], False)
  res_shape = concat([d0, num_beams, beam_length, dl])
  res = view(gather_nd(x, indices), res_shape, False)
  return res

當(dāng)然，這只是一個(gè)非常簡單的例子。如要了解更復(fù)雜的示例，請參見束搜索實(shí)現(xiàn)。借助為 ReDrafter 添加的新功能，就可以改進(jìn) TensorRT-LLM 中的 Medusa 實(shí)現(xiàn)，從而進(jìn)一步提高其性能。

ReDrafter

在 TensorRT-LLM 中的性能

根據(jù)蘋果公司的基準(zhǔn)測試，在采用 TP8 的 NVIDIA GPU 上使用 TensorRT-LLM 的 ReDrafter 最多可將吞吐量提高至基礎(chǔ) LLM 的 2.7 倍。

請注意，任何推測解碼技術(shù)的性能提升幅度都會(huì)受到諸多因素的大幅影響，包括：

GPU 利用率：推測解碼通常用于低流量場景，由于批量較小，GPU 資源的利用率通常較低。

平均接受率：由于推測解碼必須執(zhí)行額外的計(jì)算，而其中很大一部分計(jì)算最終會(huì)在驗(yàn)證后被浪費(fèi)，因此每個(gè)解碼步驟的延遲都會(huì)增加。所以要想通過推測解碼獲得任何性能上的優(yōu)勢，平均接受率必須高到足以彌補(bǔ)增加的延遲。這受到束數(shù)量、束長度和束搜索本身質(zhì)量（受訓(xùn)練數(shù)據(jù)影響）的影響。

任務(wù)：在某些任務(wù)（例如代碼完成）中預(yù)測未來的 token 更容易，使得接受率更高，性能也會(huì)因此而提升。

總結(jié)

NVIDIA 與蘋果公司的合作讓 TensorRT-LLM 變得更加強(qiáng)大和靈活，使 LLM 社區(qū)能夠創(chuàng)造出更加復(fù)雜的模型并通過 TensorRT-LLM 輕松部署，從而在 NVIDIA GPU 上實(shí)現(xiàn)無與倫比的性能。這些新特性帶來了令人興奮的可能性，我們熱切期待著社區(qū)使用 TensorRT-LLM 功能開發(fā)出新一代先進(jìn)模型，進(jìn)一步改進(jìn) LLM 工作負(fù)載。