一、大語言模型概要介紹

首先來看一下大語言模型的結構。在 Transformer 結構下的大語言模型推理的過程中，一個 token 或者一個字的生成的過程大致上可以分成兩步：

• Step 1: 根據已有信息，也就是 input 的已知信息，估計下一個 token 的概率分布。
• Step 2: 根據采樣的策略，從概率分布里面挑出最有可能的下一個 token。

這個過程有可能是以概率最大的，偏 greedy 的方式來做，要考慮到后期生成的 token 的概率，從總體上去做采樣。這是跟傳統深度學習推理不太一樣的地方。這兩步是一個循環的過程，當生成了下一個 token 之后，這個 token 會進到下一步 Step 1 里面去再生成再下一個 token，這就是推理的基本邏輯。

這里引出一個經常提到的概念，KVCache。剛才提到的在 step 1 的時候，是根據已有信息，這里的已有信息包含兩個意思，一個是原始的輸入，另一個是之后生成的 token。如果我們把一個生成的過程拆開，前面部分是最原始的輸入，生成第一個 token A。第二步從概率分布的邏輯上來說其實是要把前面的部分再加上 A 去估計下一個 token，依此循環。這會導致一個問題，計算量是在不停增長的，而且會與前面已生成的部分和 input 部分成正比，可以想象到這樣的邏輯一定會越跑越慢。

在 Transformer 結構里面，存在一個計算的特性，當前 token 的結果只與前面的 token 有關，可以把前面 token 的計算結果進行緩存，形成兩個階段：

• Prefill 階段：輸入后走一遍全部的過程，這是全量的走模型的過程，走完之后，會產生一些中間結果。這些中間結果被緩存起來，放入到圖中標紅的下一步的過程中，KVCache 在進入 attention 之前，跟現有的新生成的 token 的結果做一個 concat，然后再做計算。之后又是一個 token 生成的過程。
• Decoding 階段：通過 KVCache 的優化，decoding 階段的計算量和前面的 token 數就變得無關了。這里其實是一個近似的無關。因為在其他主要的部分都是無關的，但是在 attention 計算的地方，是被恢復成了一個全長的 token，然后進行 attention。

這就是 KVcache 的存在的意義，讓 decoding 階段的計算盡量復用以前已經計算過的結果，這樣對前部分的數量就沒有依賴，從而提高整體的推理速度。

這里存在的問題是，首先，開始輸入很多 input，之后每次只會輸入一個，比如這里輸進去的是一本書，可能是成百上千萬的輸入，之后計算的數全部都是1。所以在推理的時候就會出現這樣的一種現象，比如一張 A100 的卡，它能夠并行推理的 token 數跟 GPU 的 TFLOPS 的關系是，當 token 數增長的時候，GPU 會被更充分地利用起來，到達一個閾值之后，可能跟算值本身的時限有關，基本上到達了該 GPU 卡能夠提供的最大 TFLOPS。

我們發現在 prefill 的運轉區間，可以把 GPU 壓滿。但是因為后續 decoding 的并行每一次只預測了一個 token 數，也就是并行度非常小，在生成過程中，GPU 都處于一種不飽和的工作狀態。對于不飽和，最簡單的處理方式就是做 batch，把 batch 加大。這里面存在一個問題，如何才能把 batch size 加大？由于 KVCache 的存在，而且由于在大模型的情況下，KV cache 占用顯存非常厲害，batch size 會受到顯存的限制。

我們需要看一看，顯存到底是被怎么消耗掉的。一個正常的執行過程，包括了 prefill 和 decoding 兩個階段。一個模型加載之后會占用一部分顯存，之后第一把執行 input token 的時候，顯存會有個快速的消耗，之后隨著 token 的逐步生成，顯存的消耗在慢慢變大，這和常規的深度學習的推理非常不一樣。因為它有 prefill 過程和后面的生成的過程。這個消耗過程其一是跟長度有關系，其二是跟我們加的輸入的長度也有關系，當我們 batch size 擴大的時候，這里的顯存就會成倍地上漲。

從顯存角度來看待的話，可以列個很簡單的公式，首先是模型占用了多少參數，然后在模型的推理過程中有很多的中間變量其實也會占用一部分參數，另外有多少 token 的 KVCache 緩存也會占一部分，最終是要小于顯存的大小。

正常而言，比如一個 llama-13B 的模型，只要這種模型一上去，顯存基本就固定了。如果我們想優化的是把 batch size 做大，就可以通過優化 β，將顯存使用變小。這里主要涉及到 KVCache 的量化技術。

占用了 KV cache 的這些 token 的數量跟什么相關呢？可以這樣去列一個公式，實際上是 batch 內平均的 token 的數量。比如現在已經輸入了多少 token，以及生成了多少 token 的一個總數乘以 batch size。這里還存在一個 γ 系數，其實就是 batch 之中不同的請求之間的 token 可復用的 KVCache。

一念 LLM 框架的名字是取自一念三千，現在的大模型轉瞬之間就會生成各種不同的結果，而一念在佛教里面就是一剎那的意思。一念三千也代表我們的目標是希望大模型在一瞬之間生成世間萬象。

因此這里會對 latency 和 throughput 兩個方向重點優化。

二、一念 LLM 基本框架

一念 LLM 的基本框架如上圖所示。從上往下分別為：

最上層就是咱們經常實現的如 Llama、Baichuan、QWen 等模型，與常規的深度學習模型的推理方案不一樣，我們采用的是手寫模型。我們拋棄了計算圖，計算圖以前都是用角度算值，拼成模型，這最大的好處是有很多的靈活性，便于算法人員去實現。而這樣會帶來另外一個問題，即深度優化困難，以至于最后大家會走向一個方向，去做融合算子，然后把融合算子放到圖里面去，如果符合圖的 pattern，用融合算子去替換。

對于 Transformer 結構，因為結構非常簡單，并且現在結構基本已開始收斂，大家不再卷模型結構，更多的是卷模型效果。也就是說，這一塊我們需要實現的模型的類型是比較少的，這就讓手寫模型和手寫算子變得有利可圖。

英偉達、Facebook 等各個大廠為大語言模型場景寫的算子都非常的大。甚至像 attention 這種，按現在基本就是一個標準的，用 Flash Attention 做一個大算子來做。

我們為什么要徹底丟掉計算圖呢？因為我們還希望去優化整個模型推理過程中的顯存，只要已經規定好了模型，就可以盯著這個模型，仔仔細細地把里面顯存的使用全部去調好，使得整個顯存使用最小化。省出來的顯存都可以拿去做其他事，去做 KVCache 相關的事。

另外一個是高效調度以提高吞吐，后文中還會詳細介紹。

下面就是算子擇優，我們的底層算子，很多時候會是開源的封裝。因為其實當模型結構相對固定之后，硬件廠商會專門針對這些大的算子進行優化，提高性能，最終目的是賣掉他們的卡。現在各大廠商甚至為了不同的 tensor 的大小，去寫不同的算子，從而獲得收益的。

多硬件的支持方面，現在業界的主流框架，基本上都是英偉達，當然也有支持英特爾的 CPU。然而像支持 GPU 和華為昇騰等國產卡，還并不完善。從國內廠商的角度來看都會面臨一個問題就是高性能的 GPU 卡進不來。從業務安全的角度，我們也必須要去支持不同的硬件。

但是為什么不是按不同廠商使用不同框架，比如英偉達用一個框架，華為用華為的框架，兩邊的都能用到一個好的收益？實際上，當你在做這一層優化的時候，你會發現面臨不同的框架，需要重復做。另外各個框架本身對后期業務邏輯的適配也會不同。

我們通過上面統一框架，下面支持多種硬件，這樣就相當于做到調度這層一次優化在所有硬件上都可以用，這樣更有利于整個平臺的運營。

三、一念 LLM 框架調度

第一個問題是 ContinuousBatching 和 PageAttention，從最原始的公式來看它其實就是在有效地優化 batch size。

正常情況下，我們會把不同的請求打成一個 batch 去做 GPU 的推理，這個過程往往是一個 batch 一個 batch 地進去，要等到這個 batch 里面的請求全部處理完才退出。從 GPU 的任務調度上來說這是最簡單的使用方式，但這種使用方式最大的問題在于它的有效 batch 會越來越低，因為每個請求的輸入長度不一樣，輸出長度不一樣。比如有可能第一個請求是一個摘要的任務，會丟入一個 1,000 字的文章讓大模型用一句話總結出來；第二個請求可能是一個擴寫的任務，給了一小段話，讓模型把它擴寫成一篇長文。也就說輸入輸出的不匹配，會導致有些請求很快就結束了，有些請求還要跑很久。這樣的話，要等到所有的請求都完成，這個 batch 才會退出，這會導致有效的 batch 到后面越來越小。另外一個問題是，后面很多請求結束了，GPU 算力就會閑置。

因此很容易想到的是能否當一個請求處理完成后，就把另外新的請求輸入進去，這就是 ContinuousBatching 的基本思想。這個想法其實早在 22 年就已經被微軟提出來，但是這么好的一個想法一直都沒爆發，是因為存在一個問題，就是它的 KVCache 的顯存的操作成本比較高，比較麻煩。去年由伯克利提出的 PageAttention解決了這一問題，于是整個 ContinuousBatching+PageAttention 的機制就迎來了一個爆發期，現在成為了大語言模型推理框架的一個標配功能。

簡單講因為整塊操作成本高，所以將它切小，它采用了常規操作系統里面的內存分頁的機制來管理顯存，把 KVCache 切成不同的塊，用指針去引用的方式來進行計算。這樣就顯著降低了顯存操作的粒度。這樣 ContinuousBatching 的整體操作效率就得到了提升。

這里面還遺留了一個問題，就是 input 在請求與請求之間的共享問題，這在以前的深度學習的推理里面很少關注到，但是在推薦里面從很早的時候就已經有類似的工作。比如一個用戶多個 item 的 rank 操作，因為用戶都是一樣的，所以這一部分的推理只需要做一次。然后不同的 item 的推理多推多次，再融合到一起，再去做后面的推理工作。

請求是有共性的，這些共性的請求其實可以只算一次，然后把計算結果緩存，就可以把 KVCache 存起來，等到下一個請求，只需要處理后面的。Batch 之間也可以繼續復用，這樣整個后期請求的推理響應就可以一下提上去。

KVCache 機制看起來非常美好，但也是有成本的，因為 KVCache 會占用顯存，而且一旦緩存大，就會面臨 cache 換入換出和命中問題。比如放兩個前序在顯存里面，就得占兩份顯存。最終在具體的執行節點上面，命中率就決定了這個機制最終的收益。所以 prefix catching 更多就相當于上面公式里面的 γ。

我們在服務節點的前置加了一個叫 prefix-token 的路由器。在這里，路由需要平衡兩件事情，命中率以及傳統路由的問題，包括負載平衡、容災等問題。例如，如果前序都是一樣的，就直接打到一個節點上去，它一定會被命中。但是實際上還是會面臨負載平衡和容災的事情怎么解決的問題。所以我們構建了一個路由表，例如經常看到的一些角色扮演的服務，比如正在跟宋江聊，首先需要告訴模型你現在正在扮演的是宋江，宋江是一個什么樣的人、之前的生平、有什么樣的能力、他的性格等一些重要的事件，就像用戶簡歷一樣。然后再說你跟宋江之前聊過什么，下面請生成你準備作為宋江回復給用戶的信息。這個過程里面有大量的信息其實是跟用戶 profile 一樣。還有另外一種，比如作為愛因斯坦或者牛頓這樣的角色，對不同的角色，在進到模型之前，我們就會把前面的這一段做一個 prefix-token。對相同的一段給他指定一個具體的路由表，在有限的機器里面去選中一個機器集合。

另外剛才也提到一個問題，我們需要解決不能有太多份的 cache，cache 份數多了，顯存里就全是 cache，沒有顯存去計算當前要執行的請求了。我們通過另一個維度的管理，對于單一的節點，相當于是一個 server-set 其實是有限的。從這張圖可以感覺到，最后每個節點只命中了兩個 set，從而達到對于單個節點而言 cache 份數是相對可控的，同時又能夠在 set 的維度完成負載平衡和容災。

最后再講一下，CPU 跟 GPU 的混合推理，其實就是優化 M。

前面提到把模型從 FP16 變成 INT8 或者 INT4，這種是量化操作，效果是有損的。從框架層面需要支持，待業務評估是否使用。這里講的是無損的一些方式，計算密度的問題。我們一般講的大語言模型是計算強密集的，通常指的是 Transformer 結構部分。實際上它還有一個 token embedding 部分，這部分就是查表，類似推薦里面的 sparse 操作。往往這部分又是放在 GPU 上做的，這個表一旦大了就會占顯存。而且可以看到一般在業務應用的時候，通常都會擴 token，擴詞表。因為原始的模型里面的那些詞并不能完全覆蓋到我們的業務里。業務里面有可能會有一些特殊場景的詞。開源模型如標準的 Llama-13B 在 v2 的時候是 3.2 萬的詞表。詞表會隨業務擴展，Llama-13B 原生的詞表可能只占整個參數的百分之一。但是當把它擴到三十萬的詞表時，它對整個模型參數量的占比就會到 11.8% 的顯存。而它又是個 sparse 操作，很直接的一個想法，就是把它丟到 CPU 上去改。我們自己在測試 30 萬詞表的 Llama-13B 時，能夠有 10% 的性能提升。但這 10% 也不一定是能拿到的收益，跟詞表大小有直接相關。因為這會涉及到 CPU 跟 GPU 的聯合推理的問題，意味著 CPU 執行完的結果要拷貝到 GPU 上，多了一個成本。如果節省的顯存不足以去 cover 成本，收益就可能是負的。所以需要根據實際業務所用的模型去調整是否選用這個機制。

四、一念 LLM 在 GR 模型的應用

目前大家都在考慮大語言模型的 Transformer，生成式的結構，能否用在推薦系統中。推薦系統的推理因為量非常大，耗時要求非常高，推理成本比常規的 AI 推理要高很多。

當我們把常規的模型結構，變成一個大語言模型的結構，采用一個長序列輸入，計算量可能是成千上萬倍，成本也可能隨之增加。

生成式推薦其實是基于歷史序列去預測候選 item 的 action。這里單個用戶會有大量的 item 要預測，因為 rank 往往是千級別的。對于同一用戶的歷史序列，其實是一樣的，這是很好的符合 prefix catching 的場景。可以對這個用戶所有的 item 預測請求，做一次 prefix-cache，之后只做 item 部分的推理。這可以實現原來整個計算量是 prefix-token 加上一個待預測的 item，需要推理的 token 量乘以 item 的數量。相當于每一個 token 的計算量是乘的關系，所以會導致我們的計算量成千上萬倍的增長，因為 token 是成千上萬的。通過這個功能，就可以實現 item 變成了 item 的數量加上 token 數量，也就是把乘變成了加。這樣最后的計算量就是跟 item 數量線性相關，就跟現在正常推薦的推理相似。

這里只是解決了一個計算量的問題，還有另一個問題是 latency 的問題。如果是以萬級的 token 輸入，想要最后控制在 10 毫秒以下也是非常困難的，哪怕業務能夠接受更長時間的 latency，也不是將閾值從 10 放松到 50 毫秒這樣的一個狀態，而是要放松到秒級。這里對于 item 的預測是可以分開的。在不知道 item 的情況下，就已經知道了用戶的序列，可以提前計算。比如在用戶請求剛剛過來時，就可以把序列發給用戶了，然后等到把 item 做了召回初排之后，再去執行這一部分。這部分就只有一個 item 的 token 耗時，這也是我們傳統意義上講的 rank 的請求耗時。這個耗時就可以做到只跟 token 的最后 item 有關而跟前面的 prefix-token 的數量無關。這樣的話就可以把整個系統跟現有推薦的推理系統基本對齊。

無、未來規劃

未來的規劃就是圍繞整個架構的幾層：

1. 對模型的支持

常用的大語言模型；業務的 GR 的推薦模型的支持。

2. 調度層面的優化

計算/顯存的流水線，包括現在熱門的投機解碼等業界先進技術，會持續跟進。

3. 硬件的支持

不只是在硬件上跑起來，更重要的是在硬件上定制化算子的開發。例如華為等國內其他公司也在做加速類的硬件。要單獨提一下 CPU，因為現在最新的英特爾芯片，已經在 CPU 盒里面加進了矩陣計算的硬件單元。這種情況下，CPU 其實是可以去承擔一定程度的高密度的矩陣計算的，性能會好很多。

六、Q&A

Q1：之前在做 CTR 推理的時候做了很多類似顯存分配、動態 batch、多流并行、kernel launch 等工作，未來有哪些 CTR 推理的能力和經驗是可以借鑒到 LLM 推理上的？CTR 和 LLM 之間的區別和優化側重點有哪些共性和不同點？

A1：這里沒有提到傳統的深度學習推理里面的那些優化，準確來說這些優化全部都有效。只是在大語言模型推理場景下面，因為長序列，序列輸入是以 token 方式去做并行這樣一個特殊性，引入了一些特殊的優化方法，包括動態 batch，其實很大程度上也會是跟剛才提到的 continuous batching 結構類似，實際上 continuous batching 就是更細化的動態 batch 操作。

多流的并形其實可以用在單個請求單個 batch 前向推理的過程優化里面。這部分相當于已經有一定的 batch 了，要去生成一個 token，就要經過圖的過程。所有以前用的優化都可以繼續使用。

只是可能有一部分，比如手寫算子，圖優化，因為沒有圖了，所以也就不需要圖優化了。

簡而言之，目前的 GR 模型其實并不是一個連續生成的模型，其實對 KVCache 連續生成這一點上的依賴沒那么重。就像在現有體系下是計算圖去實現。走一次前項，再走一次前項，只不過多了一個 KVCache 的輸入。然后圖存在一些變化的部分，用傳統計算圖的一些優化方式去實現也都是沒問題的。

這個理論可能會涉及到極致優化的問題，因為手寫算子極致優化這件事情本身就是相當于損失易用性來實現的。這張圖上的應用，是跟咱們算法同學配合上線的。在大語言模型的場景下面，如果有一個新模型，就又得去支持它。從工程層面上來說，需要重新把從訓練到推理的這一套，全部的給重新弄一遍，這其實就是取決于模型結構本身的穩定性的問題。因為在大語言模型場景，結構已經非常穩定，跟推薦場景比起來簡單太多了。推薦場景會有很多的 sparse，小算子來回拼來拼去的事情。在大語言模型這種情況下全都沒有了，只剩下大算子。

像 kernel launch 這種類型的優化，在現有的大語言模型場景下？全部都可以用。現在的主流的大語言模型推理框架基本上都不是用圖。除了 tensorRT LLM 算是用圖的，但其實它里面也是很多大算子，只是用圖大致串了一下幾個大算子。

Q2：一念現在用的是連續的 batch，在大 batch 中間，模型的 forward 部分和解碼采樣是會有一個串形的執行流水線，這樣是不是可能會出現 GPU 的空隙？

A2：實際上在解碼采樣這件事情上，目前主要的優化手段還是把解碼采樣的時間段盡量放短，因為現在主要的方式仍然是 token 依次生成。當前大語言模型，因為顯存的問題，顯存占用太大了，比較難像以前一樣。這個 batch 跟那個 batch 交換著做一個流水線。

現在盡可能將一個 batch 打大，打大之后顯存就已經沒了，沒法再切到另一個 batch 的顯存來做下一步的推理，所以現在這一塊基本上是串行。現在主要的優化是想辦法把解碼采樣這一環節壓縮，把它壓的更短，不要就那個地方的解碼采樣，現在很多都在最后全部采用 GPU 算子來做，而不是在 python 層面去寫，寫了再用。其實都是為了極致的壓縮，因為它這整個過程現在就是個串行的，比以前的優化的空間小很多，不能夠讓流水線讓下一個 batch 的計算能夠進來，還需等前一個解碼采樣算完，希望盡可能的用 GPU 去算解碼采樣。

現在主流的方案基本都是這樣的，因為中間提到的 decoding 是有狀態的，其中的 KVCache 不可能在這么短的時間內（一個采樣的時間）把它換出去，再換另外一個 batch 進來。這樣的操作會得不償失。

Q3：一念有沒有做過跟 TensorRT-llm 的對比，是否有跟 A800 或者 4090 做推理。它的性價比如何？

A3：A800 或者 A100 的有測，其他普通的非線上卡沒有測。跟 tensorrt-llm 的對比的話，在具體場景，有 10%-20% 的收益。這個收益主要源于我們在下面對開源算子進行了封裝，包括 FastTransformer 的算子，vllm 的算子和 TensortRT 的算子。我們集成了開源的大語言模型推理框架，以及為業務定制的算子。這里面存在一個問題，因為做了大量的融合算子都用 FP16 在推理，在這種情況下 FP16 有精度損失。大家都知道在業務實際應用的時候，有些業務可能會有非常強的效果一致性的需求，有時候就需要把算子退化到比如 pytorch 的算子來跟訓練做對齊。

對于 TensorLRT-LLM 來說，它其實完全用 FP16 的性能并不是那么好，只有開啟 INT8 和 FP16 的混合量化的機制之后性能才上來。

Q4：如果是對于 GR 生成推理的話是更適合做一個 multi stream 的并發，還是也像 llm 那樣做一個異步大 batch 之后，再去做一個串行執行？

A4：在大語言模型情況下為什么沒有做這件事情，是因為顯存不夠。但是在 GR 的場景，模型大小不是很大，像現在做的大語言模型，13B 的規模那就是 26G 的顯存沒了。但如果模型可能只有 G 這種級別，剩下的顯存就是足夠大的，就有空間去做多 batch。對于現在大語言模型的場景，很多問題都暴露在了顯存大小這件事情上。