FlashInfer：Cascade Inference 与 Page Table

谁该读这一篇？ 想搞清 SGLang 性能的核心—attention kernel—到底在做什么的工程师。 前置阅读： 03-code-walkthrough/06-attention-backends.md、02-core-concepts/05-token-attention.md。 耗时： 35 分钟 学完能： 1. 解释 Page Table API 的三个张量及其物理意义； 2. 写出 Cascade Inference 的合并公式； 3. 分辨"FlashAttention vs FlashInfer"在 SGLang 场景下的取舍； 4. 排查 "FlashInfer wrapper plan 失败" 类问题； 5. 知道 FlashInfer 的 workspace_buffer 怎么估算。

1. FlashInfer 项目背景

FlashInfer 是 CMU/UW 团队 2024 起开源的 attention engine，专门为 LLM serving 优化：

• 可定制：用 JIT 模板生成 kernel，支持各种 attention 变体（MHA / GQA / MQA / MLA / RoPE / sliding window / softcap）。
• Page Table API：一等公民支持"非连续 KV"。
• Cascade Inference：原生支持多请求共享 prefix 的 attention。
• CUDA Graph 友好：所有元数据都做成 device tensor，wrapper 可以被 graph 录入。

SGLang 从 v0.2 起默认 attention backend。两个项目同源（LMSYS 团队同步贡献）。

2. Page Table API

FlashInfer 期望的 KV layout：

全局物理 page pool (kv_data):
  shape = [num_pages, 2 (K/V), num_layers, page_size, num_kv_heads, head_dim]

每请求的 page table:
  page_indices[r] = [p_0, p_1, ..., p_{K_r-1}]   # 这个请求用到的 page id
  page_indptr     = [start_0, start_1, ..., start_N]  # 每请求 page_indices 的起止
  last_page_len[r] = 实际占用第 K_r 个 page 多少 token

SGLang 的转换发生在 flashinfer_backend.py:953 FlashInferIndicesUpdaterDecode 里：

def update_metadata(self, forward_batch):
    # 从 forward_batch.req_to_token 翻译
    page_indices_per_req = self.req_to_token[req_pool_indices]
    page_indices = torch.cat([per_req_indices for per_req_indices in page_indices_per_req])
    page_indptr = compute_cumsum_indptr(...)
    last_page_len = (seq_lens % page_size).clamp(min=1)
    self.decode_wrapper.plan(
        page_indices=page_indices,
        page_indptr=page_indptr,
        last_page_len=last_page_len,
        num_qo_heads=...,
        num_kv_heads=...,
        head_dim=...,
        page_size=...,
        ...
    )

plan(...) 是 FlashInfer 的元数据预算调用——它会算 workspace 怎么用、kernel 怎么 launch，然后 run(q, k, v) 直接跑。

3. 一次 decode 的 attention 调用

# Model 层里：
q = self.q_proj(x)                            # [total_tokens, num_q_heads, head_dim]
k = self.k_proj(x)
v = self.v_proj(x)

# 把新算的 K/V 写进 page pool
forward_batch.token_to_kv_pool.set_kv_buffer(layer_id, out_cache_loc, k, v)

# 调 attention backend
output = self.attn_backend.forward(q, None, None, layer, forward_batch)
# attention backend 内部：
#   wrapper = self.decode_wrapper
#   output = wrapper.run(q, kv_data)        # kv_data 是 page pool

注意 forward(q, None, None, ...) 传 None 给 K/V——因为 K/V 已经在 page pool 里，wrapper 通过 page table 自己找。

4. Cascade Inference

4.1 数学等价

普通 attention 对每请求独立：

如果有 N 个请求共享 prefix $K_{shared}, V_{shared}$，朴素做法把 shared 部分复制 N 份。 FlashInfer 把它拆成两段算：

merge 利用 log-sum-exp 的稳定合并公式。

收益：

• K_{shared}, V_{shared} 在 GPU 内存里只存一份。
• shared 部分的 attention 计算可以批量做，N 个 Q 一次性吃同一段 K/V，L2 cache 命中率拉满。

4.2 多层级 cascade

RadixTree 不只是"shared / 不 shared"二元，可能多层共享：

root ──[prefix A: 1000 token]── 节点 X ──[prefix B: 500 token]── 节点 Y
                                  │
                                  ├── 请求 1（独立 tail 50）
                                  ├── 请求 2（独立 tail 80）
                                  └── 请求 3（独立 tail 30）
                                              │
                                              └── 节点 Z ── 请求 4（独立 tail 20）

MultiLevelCascadeAttentionWrapper 沿 tree 路径分段算，最后合并所有 logsumexp。 对深 RadixTree workload（agent fork），性能远超普通 paged attention。

5. Prefill vs Decode 的 wrapper

# 在 FlashInferAttnBackend.__init__：
self.prefill_wrappers = {
    WrapperDispatch.SHARED_PREFIX: MultiLevelCascadeAttentionWrapper(...),
    WrapperDispatch.SLIDING_WINDOW: BatchPrefillWithPagedKVCacheWrapper(...),
    ...
}
self.decode_wrappers = {...}

Prefill 阶段：

• 输入 Q 的 seqlen 可变（每请求 prompt 长度不同）。
• 用 BatchPrefillWithPagedKVCacheWrapper，能处理不同 seqlen 的 ragged batch。

Decode 阶段：

• 输入 Q 每请求 1 token，全 batch 整齐。
• 用 BatchDecodeWithPagedKVCacheWrapper，对应有 CUDA Graph 友好的 kernel。

6. Workspace Buffer 估算

启动时分配：

workspace_size = server_args.flashinfer_workspace_size or DEFAULT_SIZE  # 通常 128 MB
self.workspace_buffer = torch.empty(workspace_size, dtype=torch.uint8, device=device)

里面装什么：

• LogSumExp 临时缓冲（每请求 / 每 head 一份）。
• Partial output 缓冲（cascade 合并前的中间值）。
• 内部 indptr / lookup 表的 device copy。

对大 batch（256 路并发）+ 长 seqlen（32k）+ 多层 cascade，workspace 可能要 1-2 GB。 flashinfer wrapper plan 报 "workspace too small" 就是这块不够。

7. 启动失败的常见原因

[ERROR] flashinfer.plan(...) failed: workspace_size insufficient

解决：

--flashinfer-workspace-size 1073741824     # 1 GB

[ERROR] CUDA out of memory while capturing graph

可能：

• 太多 graph batch size 同时分配 workspace。
• 减 --cuda-graph-bs 列表。

[ERROR] FlashInfer: unsupported head_dim=72

罕见模型 head_dim 不在 FlashInfer 支持列表（{64, 128, 256, 192...}）。 切到 --attention-backend triton。

8. 性能比较（参考）

H100, Llama-3-8B, batch=64 decode, seqlen=2048:

差距 10-50%。Cascade 在高共享场景能再快 30%。

9. FlashInfer 与 FlashAttention 的关系

• FlashAttention：算法（tiled / online softmax / IO-aware），有 kernel 实现（Dao Lab 维护）。
• FlashInfer：一个 attention serving engine，集成了 FlashAttention 系列算法，加上 Page Table、Cascade、JIT 等 LLM serving 必备特性。

简单说：FlashInfer = FlashAttention 算法 + serving 时需要的所有"工程外壳"。

10. 调试 / Profiling

# 打开 FlashInfer 日志
export FLASHINFER_LOG_LEVEL=info
python -m sglang.launch_server ...

# 看选了哪些 wrapper / dispatch

NSys profile：

nsys profile -o sglang.qdrep python -m sglang.launch_server --model-path ... &
# 跑一个请求
nsys-ui sglang.qdrep

在时间轴上能看到 FlashInfer kernel 名（flashinfer::BatchPrefillWithPagedKVCache_kernel），便于定位瓶颈。

11. 小结

• FlashInfer 是 SGLang 默认 attention backend，专门为 paged + cascade 场景设计。
• Page Table API：page_indices + page_indptr + last_page_len。
• Cascade Inference：shared prefix 算一次，per-request tail 算一次，logsumexp 合并。
• Wrapper 模式：每次 forward 前 plan(...)，再 run(q, kv)。
• Workspace buffer 是固定开销，不够会报错。

12. 自检

1. Page table 的三张量分别表示什么？

2. Cascade Inference 怎么从 N 次重复 attention 变成"一次 shared + N 次 tail"？数学正确性靠什么？

3. Prefill 和 decode 用不同 wrapper 的原因？

4. Workspace 不够时第一反应是什么？

5. FlashInfer 不支持的 head_dim 时怎么办？

13. 下一步

• 02-speculative-decoding.md — 投机解码。
• 03-quantization.md — 量化。
• 04-cuda-graph.md — CUDA Graph。
• 03-code-walkthrough/06-attention-backends.md — backend 选择。
• 参考：FlashInfer 论文、FlashInfer GitHub。

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