AI Infrastructure

Block Attention 與 KV Cache 重用：RAG 場景的推理加速新途徑

透過語義分割與塊蒸餾，讓 Block Attention 達成近乎全注意力的性能，開啟 RAG 場景 KV Cache 完全重用的路徑

簡介

RAG（Retrieval-Augmented Generation）系統面臨一個根本性瓶頸：每次使用者查詢，模型都必須對所有檢索回來的文檔重新計算 KV Cache。這意味著同一份文檔被 100 個查詢引用，就會產生 100 次完全重複的計算。在生產環境中，這些重複計算佔據了 RAG 推理成本的極大比例。

來自 arXiv 2605.15913 的論文提出了一條根本不同的路徑：Block Attention。核心想法是將輸入文本自動分割為語義獨立的塊，每個塊只在自己的 Token 之間做注意力運算。這樣一來，每個文檔塊的 KV Cache 就不再依賴查詢上下文，可以被預先計算並重複使用。

但這個看似簡單的改動帶來兩大挑戰：第一，如何自動且正確地切分文本；第二，砍掉跨塊注意力之後，模型性能怎麼辦？這篇論文分別用 SemanticSeg 資料集和三組件塊蒸餾框架回答了這兩個問題，實驗結果相當令人驚訝——塊注意力下的模型性能接近全注意力。

作為 AI Infrastructure 工程師，這篇工作對 LLM 推理服務基礎設施的意義非常直接，值得深入拆解。

Block Attention 的計算優勢

傳統的全注意力機制下，每個 Token 都可以注意序列中的所有其他 Token。在 RAG 場景中，這意味著查詢和文檔之間存在完全的交叉注意力。問題是，一旦查詢改變，所有 KV 值的計算都會受到影響——因為注意力權重依賴於當前的 Q（查詢），而 Q 又是查詢上下文的函數。

Block Attention 把輸入分割成多個獨立塊，注意力只在塊內發生，不跨塊。這帶來一個關鍵性質：每個塊的 KV Cache 只依賴塊內的 Token，與外部查詢完全無關。

1
2
3
4
5
6
7
8
9

# Full Attention -- query 依賴所有 document tokens
Query:  [Q1] [Q2]
Doc:    [D1] [D2] [D3] [D4]
KV Cache for Doc = f(Q1, Q2, D1, D2, D3, D4)  # 每次查詢都要重算

# Block Attention -- 查詢塊與文檔塊各自獨立
Query Block:  [Q1] [Q2]     KV = f(Q1, Q2)
Doc Block 1:  [D1] [D2]     KV = f(D1, D2)   # 可預計算!
Doc Block 2:  [D3] [D4]     KV = f(D3, D4)   # 可預計算!

以下是不同注意力機制在 RAG 場景中的比較：

注意力模式 KV Cache 重用 跨區段資訊流動 部署複雜度
全注意力 無法重用 自由流動 低
滑動窗口（如 Mistral） 有限重用 部分重疊 中
前綴快取（Anthropic） 前綴可重用 前綴內自由 中
Block Attention 塊完全重用 僅通過匯聚 Token 高（需蒸餾）

值得注意的是，Block Attention 的塊可以任意組合、重排序、做子集選擇，不受前綴匹配的約束。這對動態檢索場景（每次查詢檢索不同子集的文檔）來說，靈活性遠超現有的 Prompt Caching 方案。

SemanticSeg：讓機器學會「在哪裡切」

Block Attention 的第一個難題是分割點。機械式按固定 token 數切分會破壞語義完整性——想像一個跨邊界的長句，前半段在一個塊、後半段在另一個塊，兩者之間又無法互相注意，資訊必然丟失。

論文提出了 SemanticSeg，首個大規模文本語義分割資料集：

屬性 數值
總實例數 30,000+
類別數 16（書籍、程式碼、網頁、對話等）
文本長度範圍 2K - 32K tokens
分割器特性 輕量級、可控粒度

這個資料集的核心價值不在於某個具體的分割器，而在於建立了一個可量化的基準。可控粒度的設計尤其實用：部署時可以根據場景在細粒度（句子級）和粗粒度（章節級）之間調整塊大小，平衡 KV Cache 顆粒度與塊間通訊成本。

細粒度切割產生更多塊，每塊的 KV Cache 更小但需要更多匯聚 Token 來傳遞跨塊資訊；粗粒度切割則相反。這是一個工程上需要針對實際部署場景調優的權衡點。

三組件塊蒸餾框架

Block Attention 的第二個難題更根本：砍掉跨塊注意力之後，模型能維持多少性能？論文的答案是透過知識蒸餾從全注意力教師轉移知識到塊注意力學生，並設計了三個關鍵組件。

塊匯聚 Token（Block Sink Tokens）

每個塊添加若干特殊的 Sink Token，塊內所有 Token 都注意這些 Token。Sink Token 累積整個塊的壓縮表示，作為塊間通訊的固定大小介面。

這個設計的巧妙之處在於：Sink Token 的數量直接控制了塊間資訊傳遞的頻寬。它類似 Perceiver 架構中的 Latent Array，但針對塊注意力做了專門優化。代詞跨塊引用、推理鏈跨塊延續等場景，其跨塊依賴的資訊被壓縮到 Sink Token 中傳遞。

塊 Dropout（Block Dropout）

訓練時以機率 p_drop 隨機丟棄整個塊。這直接模擬了生產 RAG 中的現實：檢索回來的文檔不一定都相關，有時候還包含噪聲。Block Dropout 迫使模型學會不依賴任何特定塊的存在，提升魯棒性。

這是一個看似簡單但效果顯著的正則化策略。它還有一個實用副作用：模型學會在部分塊缺失時仍能做出合理推論，這意味著部署時可以更安全地做動態塊子集選擇。

Token 級損失加權（Token-Level Loss Weighting）

標準蒸餾對所有 Token 均等計算損失，但在塊注意力場景中，不同 Token 受影響的程度差異很大。論文提出的加權策略聚焦於：

  • 塊邊界 Token：高權重——受注意力截斷影響最大
  • 匯聚 Token：高權重——承載壓縮的塊間資訊，是整個框架的通訊樞紐
  • 資訊密集 Token（實體名稱、推理關鍵步驟）：高權重
  • 填充 Token：低權重——浪費計算在無意義的位置上毫無幫助
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

# 概念性的 Token 級損失加權示意
def compute_distill_loss(student_logits, teacher_logits, token_weights):
    """
    token_weights: 根據 Token 在塊內位置和語義重要性計算
    - 塊邊界: weight = 2.0
    - 匯聚 Token: weight = 3.0
    - 關鍵實體: weight = 1.5
    - 塊內部普通 Token: weight = 1.0
    - 填充 Token: weight = 0.1
    """
    per_token_loss = kl_divergence(student_logits, teacher_logits)
    return (per_token_loss * token_weights).mean()

實驗結果與性能表現

跨多個模型和基準測試，論文報告了幾個關鍵結果：

  • 分割器在語義一致性上超越啟發式和統計基線
  • 塊蒸餾模型在塊注意力限制下達到近乎全注意力的性能
  • KV Cache 重用率接近 100%
  • 多查詢場景下推理吞吐量顯著提升

這組結果中最重要的不是某個具體數字，而是挑戰了一個根深蒂固的假設：「跨文檔注意力是不可少的」。實驗表明，透過適當的蒸餾訓練，跨塊注意力的大部分功能可以被壓縮到 Sink Token 中。如果這個結論在更多場景下成立，意味著生產 RAG 系統中跨文檔注意力的計算開銷可能是可避免的。

架構分析：與現有技術的定位 

  graph TD
    A[KV Cache Optimization Landscape] --> B[Memory Management Layer]
    A --> C[Computation Structure Layer]
    A --> D[Model Training Layer]

    B --> B1[PagedAttention - vLLM]
    B --> B2[Prefix Caching - Anthropic / Google]
    B --> B3[KV Cache Quantization - TurboQuant]

    C --> C1[Sliding Window Attention - Mistral]
    C --> C2[Block Attention - This Work]

    D --> D1[Knowledge Distillation]
    D --> D2[Block Distillation - This Work]

    C2 -.->|semantic blocks managed by| B1
    D2 -.->|enables| C2
    C2 -.->|more flexible than| B2

Block Attention 在這個技術版圖中的定位很清晰：它在計算結構層做優化，透過改變注意力的結構模式來 enable KV Cache 重用。幾個重要的互補關係：

與 PagedAttention 互補而非替代。PagedAttention 在物理記憶體層面管理 KV Cache 的分配與回收，Block Attention 在語義層面管理 Cache 的粒度與重用。塊注意力的語義塊完全可以被 PagedAttention 的記憶體管理器管理，兩者在不同層面優化，可以疊加。

比前綴快取更靈活。Anthropic 和 Google 的 Prompt Caching 只支持嚴格的前綴匹配——查詢必須以完全相同的文檔序列開頭。Block Attention 支持任意塊子集的組合，在動態檢索場景中每個查詢可能引用不同文檔子集的情況下，靈活性有質的差異。

與 KV Cache 量化互補。TurboQuant 等量化方法可以進一步壓縮每個塊的 KV Cache 体积，與 Block Attention 的重用機制疊加。

風險與工程挑戰

從工程落地角度，這個方案有幾個需要正視的問題：

多跳推理的資訊瓶頸。需要強跨文檔推理的任務（多跳 QA、跨文檔比較）可能因 Sink Token 的頻寬限制而受損。Sink Token 數量是一個需要針對場景調優的超參數——太少則資訊丟失，太多則計算量增加、抵消效率收益。

分割器品質直接決定上限。語義切分的錯誤會直接導致注意力品質下降。在生產環境中需要對分割器做持續監控，確保其在各類文本上的魯棒性。

每個基礎模型需重新蒸餾。塊蒸餾是針對特定模型的過程，換一個基礎模型就需要重複整個蒸餾流程。這增加了模型迭代時的部署成本和時間。

與 FlashAttention 的整合。塊內注意力可以使用 FlashAttention 加速，但整體框架（Sink Token 處理、塊級調度）需要額外的工程整合，目前的深度學習框架並未原生支援這種模式。

總結

Block Attention 為 RAG 系統提供了一條從「每次查詢重算所有文檔」到「文檔 KV Cache 預計算重用」的實用路徑。SemanticSeg 填補了文本語義分割的基準空白，三組件蒸餾框架則展示了「改變注意力結構但不損失性能」的可能性。

這項工作的核心啟示是：跨文檔注意力的計算開銷可能是不必要的。如果透過蒸餾就能將跨塊依賴壓縮到固定大小的匯聚 Token 中，那麼 RAG 場景的推理成本可以大幅降低。對於運行大規模 RAG 服務的基礎設施團隊來說，這可能是下一個重要的效率突破口。

當然，目前這仍處於研究階段。多跳推理的壓力測試、不同模型架構的遷移成本、生產環境的分割器監控等工程問題，都需要在實際部署中驗證。但方向是明確的：KV Cache 重用正在從物理層面的優化（PagedAttention、Quantization）向語義和結構層面拓展，Block Attention 是這個趨勢中一個值得密切追蹤的里程碑。

參考論文：Towards Generalization of Block Attention via Automatic Segmentation and Block Distillation (arXiv 2605.15913)

© 2026 · Engineering Notes