源码精读：Attention 后端

谁该读这一篇？ 想理解 SGLang 是如何把"逻辑请求 batch"转成"GPU 上一次 attention kernel"的工程师；做 perf tuning 的同学。 前置阅读： 02-core-concepts/05-token-attention.md、04-model-runner.md。 耗时： 40 分钟 学完能： 1. 列出 SGLang 支持的主要 attention backend； 2. 解释 FlashInfer 的 wrapper / page table / cascade 三个核心概念； 3. 在源码里找到 backend 注册和选择的位置； 4. 看到 attention 慢能判断是 backend 选错、kernel 不匹配还是数据不连续； 5. 知道什么时候该切换 backend。

1. Backend 概览

srt/layers/attention/ 下有 20+ 个 backend 文件，但常用的就几个：

切换：

python -m sglang.launch_server --attention-backend flashinfer
# 或 triton, flashinfer-mla, flashmla, torch, ...

2. 抽象基类

源码：base_attn_backend.py。

每个 backend 都实现：

class AttentionBackend(ABC):
    def init_forward_metadata(self, forward_batch):
        """在 forward 前预算 page table、index、wrapper 等。"""

    def forward(self, q, k, v, layer, forward_batch):
        """单层 attention 计算。"""

    def init_cuda_graph_state(self, max_bs):
        """CUDA Graph 录制时分配持久化张量。"""

    def init_forward_metadata_capture_cuda_graph(self, ...):
        """CUDA Graph 录制时调用。"""

    def init_forward_metadata_replay_cuda_graph(self, ...):
        """CUDA Graph 回放时调用。"""

注册表：attention_registry.py。 ModelRunner 启动时按 --attention-backend 字符串查注册表实例化。

3. FlashInfer Backend（114+）

源码：flashinfer_backend.py:114（class FlashInferAttnBackend）。

3.1 Wrapper 概念

FlashInfer 的设计是把"一次 attention 调用"包装成一个 wrapper 对象：

self.decode_wrapper = BatchDecodeWithPagedKVCacheWrapper(workspace_buffer, ...)
self.prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper(workspace_buffer, ...)

每个 wrapper 内部持有：

• workspace_buffer：一块预分配显存（cuBLAS workspace）。
• page_table 缓存。
• 当前 batch 的元数据。

调用 wrapper.plan(...) 配置好元数据，再 wrapper.run(q, k, v) 跑 kernel。

3.2 Page Table

FlashInfer 期望传入 "page table"：

page_indices:    [page_id_0, page_id_1, ..., page_id_K-1]
page_indptr:     [start_0, start_1, ..., start_N]    # 每请求的 page 起止
last_page_len:   [len_0, len_1, ..., len_N-1]        # 每请求最后一 page 实际占多少 token

这个布局源自 vLLM 的 PagedAttention，FlashInfer 把它做成 first-class API。 SGLang 的 RadixCache 在 init_forward_metadata 时把"每请求当前命中的 token slot 列表"翻译成上面三个张量。

3.3 Cascade Inference

最酷的特性：当多个请求共享一段 prefix 时，FlashInfer 可以一次性算"shared prefix part"，再算每请求"unique tail part"，最后合并 attention logsumexp。

数学等价但运算量大幅减少：

其中 $K_s$ / $V_s$ 是 shared prefix，$K_t$ / $V_t$ 是 unique tail。 共享部分只被读一次（K/V 命中 L2 cache），不像传统做法把整段 K/V 复制 N 份。

SGLang 用 MultiLevelCascadeAttentionWrapper 直接吃 RadixTree 的多层结构。

3.4 决定权所在

# WrapperDispatch:
class WrapperDispatch(Enum):
    SLIDING_WINDOW = "sliding_window"
    CROSS_ATTENTION = "cross_attention"
    SHARED_PREFIX = "shared_prefix"        # Cascade
    ENCODER_DECODER = "encoder_decoder"
    ...

init_forward_metadata 根据当前 batch 形态选 wrapper（flashinfer_backend.py:953 的 FlashInferIndicesUpdaterDecode）。

4. FlashInfer MLA Backend

MLA（Multi-head Latent Attention）是 DeepSeek-V2/V3 的核心创新： 把 KV "压缩"到一个低维 latent 空间，attention kernel 直接吃压缩 KV。

源码：flashinfer_mla_backend.py。 和普通 FlashInfer 的差异主要在：

• KV 张量的形状（latent_dim 而不是 head_dim × num_kv_heads）。
• 用 BatchPrefillWithRaggedKVCacheWrapper 而不是 Paged 版本（MLA 内部不分 page）。
• 与 MLATokenToKVPool 配合。

SGLang 在 DeepSeek 系上对 vLLM 有性能优势，主要靠这条路径。

5. 怎么选 backend

启动时自动选（server_args.py 默认值）的逻辑大致：

if model_is_mla:
    backend = "flashinfer_mla"
elif gpu_compute_capability >= 9.0:    # Hopper +
    backend = "flashinfer"
elif gpu_compute_capability >= 8.0:    # Ampere
    backend = "flashinfer"
else:
    backend = "triton"                  # 老 GPU

可以强制：

--attention-backend flashinfer        # 通用
--attention-backend flashinfer-mla    # DeepSeek
--attention-backend flashmla          # 另一 MLA 实现
--attention-backend triton            # 老卡

6. 一次 forward 里 attention 的具体路径

sequenceDiagram
    participant MR as ModelRunner.forward
    participant AB as AttentionBackend
    participant FI as FlashInfer wrapper
    participant K as GPU kernel

    MR->>AB: init_forward_metadata(forward_batch)
    AB->>AB: 算 page_indices / page_indptr / last_page_len
    AB->>FI: wrapper.plan(...)
    Note over FI: 预算 workspace, kv layout
    loop 每层 layer
        MR->>MR: linear Q/K/V projection
        MR->>AB: backend.forward(q, k, v, layer, forward_batch)
        AB->>FI: wrapper.run(q, k, v)
        FI->>K: launch attention kernel
        K-->>FI: output
        FI-->>AB: output
        AB-->>MR: output
    end

注意：

• init_forward_metadata 一次性算好整个 batch 所有层的 metadata（不是每层算）。
• 每层 forward 调用复用同一 wrapper，只是输入 q/k/v 不同。
• Decode 模式下 wrapper 是 decode_wrapper，prefill 是 prefill_wrapper，混合 batch 时按 token 切分。

7. Workspace buffer

FlashInfer wrapper 需要一块"草稿纸"——workspace_buffer：

self.workspace_buffer = torch.empty(
    workspace_size_bytes,        # 通常几百 MB 到几 GB
    dtype=torch.uint8,
    device=device,
)

这块显存预分配，运行时复用；它是 KV pool 之外另一块固定显存开销。 启动日志 "Memory pool end. avail mem=XX GB" 里已经扣掉这块。

碎片化导致 workspace 不够时会 OOM，这是诊断 "起 server 时 OOM" 的常见原因。

8. Triton Backend

flashattention_backend.py 和 triton_*.py 系列。

特点：

• Pure Python + Triton，调试容易、易改。
• 性能比 FlashInfer 略差（5-15%）。
• 适合：开发调试、老 GPU、FlashInfer 不支持的特殊形状。

9. 自定义 attention backend

from sglang.srt.layers.attention.base_attn_backend import AttentionBackend

class MyBackend(AttentionBackend):
    def init_forward_metadata(self, fb): ...
    def forward(self, q, k, v, layer, fb): ...
    def init_cuda_graph_state(self, max_bs): ...
    ...

然后在 attention_registry.py 注册：

register_attention_backend("my_backend", MyBackend)

启动用 --attention-backend my_backend。

用途：实验性 kernel、特定硬件（如 TPU/NPU）、新 attention 算子（DSA、NSA）。

10. 性能调优 cheatsheet

11. 关键源码索引

12. 小结

• SGLang 支持多个 attention backend；默认 FlashInfer（DeepSeek 用 MLA 变体）。
• FlashInfer = Wrapper + Page Table + Cascade，吃 SGLang 的 token-level KV 无压力。
• Cascade Inference 是 SGLang 在多分支 / 长共享 prefix 上的性能秘密武器。
• Workspace buffer 是固定显存开销；不够会启动 OOM。
• 想自定义只需实现 AttentionBackend 接口 + 注册。

13. 自检

1. FlashInfer 和 FlashAttention 的关系？SGLang 为什么选前者？

2. Page table 的三个张量分别是什么？

3. Cascade Inference 的数学等价怎么写？

4. MLA 模型为什么不能直接用普通 FlashInfer wrapper？

5. workspace buffer 太大会怎么样？太小会怎么样？

14. 下一步

• 04-optimizations/01-flashinfer.md — 深入 Cascade。
• 04-model-runner.md — 上层调用方。
• 02-core-concepts/05-token-attention.md — 设计动机。
• 源码：sglang/python/sglang/srt/layers/attention/。

上游源码：sglang/python/sglang/srt/layers/attention/。