SGLang 從入門到精通

vLLM 是本號的常客，SGLang 寫的不多，主要是我用它也不多，之前偶爾寫過 SGLang 怎么用、跑什么模型，也比較淺

最近，吳恩達 DeepLearning 上最新更新了 SGLang 底層原理短課

我跟著學了兩個課時，感覺受益匪淺，推薦給大家

deeplearning.ai/short-courses/efficient-inference-with-sglang-text-and-image-generation/

我只學了自己感興趣的 L2 和 L3，本文也算是學習筆記

為什么要學推理優化？

這門課的名字叫 Efficient Inference with SGLang，由 SGLang 的作者團隊和 DeepLearning.AI 聯合出品

說實話，大部分人用 vLLM、SGLang 部署模型，都是pip install然后一行命令啟動服務，能跑就行

但你有沒有想過：

• ? 為什么 KV Cache 能讓推理速度提升 10 倍？

• ? 為什么 SGLang 比 vLLM 在某些場景下快那么多？

• ? RadixAttention 到底是個什么東西？

這門課就是來回答這些問題的

不講廢話，直接手寫代碼，從 Attention 公式寫到 KV Cache，再到 Radix Tree，一步步把原理拆給你看，配上我自己的理解。

L2：KV Cache——從 O(n2) 到 O(n) 的質變 先搞清楚一件事：推理為什么慢？

大語言模型生成文本是一個 token 一個 token 蹦出來的（自回歸生成）

每生成一個新 token，模型都要跑一遍 Attention 機制，用當前 token 的 Query 去和所有之前 token的 Key 做點積，算出注意力權重，再加權求和所有 Value

關鍵洞察來了：每個 token 的 Key 和 Value 一旦算出來就不會變

但是如果不緩存，每生成一個新 token，模型就要把之前所有 token 的 K 和 V 重新算一遍

生成 n 個 token，總計算量是 O(n2)，這就是推理慢的根本原因。

Attention 公式，手寫一遍

課程用的是 DeepSeek-R1-Distill-Qwen 1.5B 模型，雖然小，但 Attention 架構和 70B 模型完全一樣——Grouped Query Attention（GQA），所有原理直接適用于大模型。

先看核心公式：Attention(Q, K, V) = softmax(Q·K^T / sqrt(d_k)) · V

把這個公式翻譯成了 Python 代碼：

def_attention_impl(q, k, v, scale, mask):
# 核心：softmax(Q @ K^T / sqrt(d_k)) @ V
# Q @ K^T —— 算每對 (query, key) 的注意力分數
    scores = torch.matmul(q, k.transpose(-2, -1)) * scale
# 因果遮罩——未來位置變成 -inf，softmax 后變 0
    scores = scores.masked_fill(~mask, float("-inf"))
# 歸一化為注意力權重
    probs = torch.softmax(scores, dim=-1)
# 按權重對 Value 加權求和
return torch.matmul(probs, v)

然后還有一個處理 GQA（Grouped Query Attention）的包裝函數

GQA 是 DeepSeek、Llama 等現代模型的標配——多個 Query Head 共享同一組 K/V Head，比如 64 個 Query Head 共用 8 個 K/V Head，KV Cache 直接縮小 8 倍，精度損失很小

defsimple_causal_attention(query, key, value, **kwargs):
# 支持 GQA 的因果注意力
    Dh = query.shape[-1]
    scale = 1.0 / (Dh ** 0.5)
# GQA: 多個 Query Head 共享一組 K/V Head
    gqa_group_size = query.shape[1] // key.shape[1]
    key = key.repeat_interleave(gqa_group_size, dim=1)
    value = value.repeat_interleave(gqa_group_size, dim=1)
# ...后續和標準 Attention 一樣

課程做了一件很有意思的事：用 monkey-patch 把 PyTorch 內置的F.scaled_dot_product_attention換成自己寫的版本，這樣模型每一層都跑自己的代碼

跑出來的結果和原版完全一致——token 級別一模一樣，只是慢很多（純 Python vs CUDA 內核嘛）

沒有 KV Cache vs 有 KV Cache

這是課程里最直觀的實驗

不用 KV Cache（樸素方式）：

# 每一步都從頭喂入整個序列
text_no_cache = auto_regressive_decode(
    tiny_llm, input_text,
    max_new_tokens, temperature=0.0
)
# 總計算量：sum(9, 10, 11, ..., 24) = 264 次 token 計算

9 個 prompt token，生成 16 個新 token。每一步都要把之前所有 token 重新過一遍模型

第 1 步處理 9 個 token，第 2 步處理 10 個……第 16 步處理 24 個，加起來 264 次 token 計算。

用 KV Cache（優化方式）：

# Prefill 階段一次性處理所有 prompt token，存下 K/V
# Decode 階段每步只處理 1 個新 token
text_kv_cache = auto_regressive_decode_with_kv_cache(
    tiny_llm, input_text,
    max_new_tokens, temperature=0.0
)
# 總計算量：9 (prefill) + 15 (每步 1 個) = 24 次 token 計算

同樣 9 個 prompt token + 16 個新 token

Prefill 階段一口氣處理 9 個 token，把所有 K/V 存起來

之后每步只要處理 1 個新 token，從緩存里讀之前的 K/V 就行，總計算量 24 次。

264 vs 24，計算量少了 11 倍。

實測在 1.5B 模型上大約 2 倍實際加速，序列越長差距越大——1000 個 token 的輸出，沒有 KV Cache 需要 50 萬次計算，有了 KV Cache 只需要約 1000 次

這就是"實用"和"不可用"之間的距離

而且最關鍵的——輸出完全一樣，一個 token 都不差。數學上嚴格等價，只是不重復做已經做過的工作

小結：KV Cache 的核心思想

兩個階段：

1. 1.Prefill：并行處理整個 prompt，計算并存儲所有 token 的 K、V

2. 2.Decode：逐個生成新 token，每步只計算新 token 的 Q/K/V，之前的 K/V 直接從緩存讀

一句話：*算一次，存起來，反復用。

L3：跨請求緩存——RadixAttention

KV Cache 解決了單個請求內的重復計算問題，但有一個更扎心的問題：

兩個用戶問了同一篇文檔的不同問題，模型對同一篇文檔的 KV 算了兩遍，這是不是浪費？

答案是：當然是

在 RAG 場景里，一個 prompt 可能 90% 是文檔內容（幾百個 token），只有 10% 是用戶問題（幾十個 token）

100 個用戶對同一篇文檔提了 100 個問題，那就是 100 次完全相同的 KV 計算

幾萬個 token 白算了

SGLang 的 RadixAttention 就是來解決這個問題的

Radix Tree（基數樹）是什么

簡單說，Radix Tree 就是一個按 token 序列索引 KV Cache 的樹形數據結構。

課程里實現了一個簡化版的FlatRadixTree：

classCacheEntry:
# 把 token 序列和對應的 KV Cache 配對存儲
def__init__(self, token_ids, kv_cache):
self.token_ids = list(token_ids)
self.kv_cache = kv_cache

 classFlatRadixTree:
# 簡化版基數樹（線性掃描，便于理解）
def__init__(self):
self.entries = []

 definsert(self, token_ids, kv_cache):
self.entries.append(CacheEntry(token_ids, kv_cache))

 defsearch(self, token_ids):
# 找最長匹配前綴
        best_match, best_len = None, 0
for entry inself.entries:
            match_len = 0
for a, b inzip(entry.token_ids, token_ids):
if a != b:
break
                match_len += 1
if match_len > best_len:
                best_len = match_len
                best_match = entry
return best_match

生產級 SGLang 用的是 O(log n) 的查找，課程用線性掃描，原理一模一樣，就是方便你看懂

四步流程

每個請求進來，經過四步：

1. 1.Traverse（遍歷）：radix.search(token_ids)在樹中查找最長匹配前綴

2. 2.Reuse（復用）：如果匹配到了，直接加載已緩存的 KV 張量

3. 3.Compute（計算）：只對未匹配的后綴（比如用戶的新問題）運行模型

4. 4.Store（存儲）：radix.insert()把新算出來的 KV 存回樹里

用代碼看更清楚：

radix = FlatRadixTree()  # 空樹

 for question in article_questions:
    prompt = construct_prompt(article, question)
    token_ids = tiny_llm.tokenize(prompt)

 # Step 1: 搜索最長匹配前綴
    prefix_cache = radix.search(token_ids)

 # Steps 2 & 3: 復用緩存的 KV，只計算后綴
    output, cached_req = tiny_llm.generate_with_prefix_cache(
        prompt,
        max_new_tokens=16,
        prefix_cache=prefix_cache,
        temperature=0,
    )

 # Step 4: 存回樹里，后續請求受益
    radix.insert(cached_req.token_ids, cached_req.kv_cache)

作為實驗，準備了兩篇 SGLang 技術文章（各約 2000 字符），每篇 6 個問題

不用 prefix caching：6 個問題每個都要從頭處理整個文檔 + 問題，耗時基本一樣

用 RadixAttention：Q1 是冷啟動（cache miss），跟不緩存一樣慢。但 Q2 到 Q6 直接命中緩存——文檔部分（約 97% 的 token）全部跳過，只處理問題部分的那幾十個 token

實測大約2 倍加速，總共省了約 20 秒

你可能覺得 2 倍不夠震撼？那是因為實驗用的文章比較短

在生產環境里，一個 RAG prompt 可能有 2000 個 token 的文檔 + 10 個 token 的問題。如果 90% 的 token 都命中緩存，只需要計算 10%，那就是10 倍 Prefill 加速

混合文檔負載實驗

真實生產環境不會乖乖按文檔分組——用戶請求是隨機到達的

第三個實驗：兩篇文章的 12 個問題隨機打亂順序

# 12 個 prompt，隨機混合兩篇文章的問題
random.seed(42)
random.shuffle(all_prompts)


 radix_multi = FlatRadixTree()
for tag, article_name, prompt in all_prompts:
    token_ids = tiny_llm.tokenize(prompt)
    prefix_cache = radix_multi.search(token_ids)
# ... 生成 + 存儲

結果：只有 2 次 cache miss（每篇文章的第一次請求），剩下 10 次全部命中

命中率 83%

這就是 Radix Tree 和普通單條緩存的區別——樹可以同時維護多個分支，切換文檔不會把另一篇的緩存踢掉

每個請求獨立匹配自己的前綴，互不干擾

隨著流量增長，2 次冷啟動被上千個請求分攤，平均延遲趨近于 cache hit 的延遲

這就是為什么 SGLang 在生產環境里這么快

適用場景一覽

場景

緩存了什么

典型加速

RAG 系統

文檔上下文

5-10x

聊天機器人

System Prompt + 對話歷史

3-5x

Few-Shot 學習

示例樣本

4-8x

代碼生成

倉庫上下文

3-6x

核心邏輯都一樣：共享前綴越長，加速越大。

兩節課串起來

L2 和 L3 解決的問題其實是一個遞進關系：

• ?L2 的 KV Cache：解決單個請求內的重復計算，同一個請求里，已經算過的 token 的 K/V 不用再算

• ?L3 的 RadixAttention：解決跨請求的重復計算，不同請求共享相同前綴時，K/V 只算一次

兩層疊加，效果是乘法級的——單請求內不浪費，跨請求也不浪費

代碼模式簡單到令人發指——三行搞定：

prefix_cache = radix.search(token_ids)    # 搜
output = model.generate(prompt, prefix_cache=prefix_cache)  # 用
radix.insert(token_ids, kv_cache)          # 存

說實話，之前用 SGLang 就是跑跑 benchmark，知道它快，但不知道為什么快

這門課最大的價值是讓你親手把 KV Cache 和 Radix Tree 寫一遍——寫完之后，你看 SGLang 的源碼就不再是天書了

推薦給兩類人：

1. 1.想搞推理優化的工程師：這是入門基礎，必須搞懂

2. 2.用 vLLM/SGLang 部署模型的人：知道底層原理，遇到性能問題才知道往哪里調

課程后面還有 L4（SGLang Diffusion，把 caching 思想用到圖像生成）和 L5（SGLang Router，多引擎路由），等我學完再寫。

制作不易，如果這篇文章覺得對你有用，可否點個關注。給我個三連擊：點贊、轉發和在看。若可以再給我加個，謝謝你看我的文章，我們下篇再見！