FlashInfer 原理：为 LLM serving 设计的 attention 引擎

① 一句话和全景图

FlashInfer 是一个“attention 代码生成器 + 运行时调度器”：编译期根据 attention 变体、任务类型和 KV 存储格式，用 JIT 生成专门的 kernel；运行期根据这一步每个请求的长度，在 CPU 上算出负载均衡的调度计划，再让 GPU 按计划执行。

② 背景：serving 里的 attention 和训练时不一样

FlashAttention（T5 写过）解决了“不把 S×S 矩阵写进显存”的问题，它是为训练设计的：一个 batch 里的序列长度通常一样，Q、K、V 都是连续的稠密张量，attention 只有一种标准形式。

LLM serving 的情况完全不同：

论文的出发点就是：每个 serving 系统都在为这些情况的某个子集手写 attention kernel，维护成本高，而且常常不是最优的。FlashInfer 想用一个统一的引擎覆盖它们。

③ 四个观察，对应四组设计

观察 4 的量化：FlashAttention 的算术强度大约是 O(1 / (1/l_qo + 1/l_kv))。serving 里 l_qo ≤ l_kv，所以强度 ≈ O(l_qo)：prefill 时 l_qo 是几千，是 compute-bound 的；decode 时 l_qo = 1，是 memory-bound 的（L2 讲过）。GQA 把强度提高到 O(g · l_qo)，g 是每组的 query head 数。同一套 kernel 配置不可能同时适合这两端。

④ 基石：attention state 和合并运算 ⊕

对一个 query、一部分 key（下标集合 I），定义它的 attention state = (O(I), LSE(I))：O 是只看这部分 key 时的输出，LSE 是这部分分数的 log-sum-exp。两个不相交部分的 state 可以精确合并：

O(I ∪ J)   = ( e^{LSE(I)}·O(I) + e^{LSE(J)}·O(J) ) / ( e^{LSE(I)} + e^{LSE(J)} )
LSE(I ∪ J) = log( e^{LSE(I)} + e^{LSE(J)} )

这个 ⊕ 满足交换律和结合律，所以任意多块可以按任意顺序、任意分组合并。论文的原话是：在 FlashInfer 里，attention state 是 attention 运算的标准输出，⊕ 是标准的归约运算，地位相当于 GEMM 里的加法。

源码：flashinferstate_t，合并在 flashinferstate_t::merge；独立的合并 kernel 在 flashinferMergeStateKernel。

⑤ 设计 1：统一的 block-sparse KV 存储

观察：所有的 KV 存储格式其实都是一个稀疏矩阵

把 attention 看成一个矩阵：行是 query，列是 KV cache 里的所有位置。第 i 行第 j 列“有值”，表示 query i 要看 KV 位置 j。

• 分页 KV（vLLM 的 PagedAttention）：每个请求只看自己的那几页。
• radix tree（SGLang）：前缀共享的请求看同一批节点。
• 树形 attention（speculative decoding）、按重要性挑选的 KV（Quest 之类）：每个 query 看的位置更不规则。

这些都可以用同一种格式描述：BSR（Block Sparse Row，块稀疏行）。矩阵被切成 B_r × B_c 的块，只存非零块，用 indptr（每行块从哪开始）和 indices（每个非零块是第几列块）描述。

三个 decode 请求（每个 1 个 query），page 大小为 B_c，物理页共 8 个：A 用页 5、1、6，B 用页 2，C 用页 7、3。

每个彩色格子是一个非零块（B_r = 1 个 query × B_c = 一页的 token）。圈里的数字是这一页在请求自己序列里的顺序。用 BSR 描述：indptr = [0, 3, 4, 6]，indices = [5, 1, 6, 2, 7, 3]，这正是 F1 里的 CSR 页表。

论文的定义：query 和输出用不带 padding 的 ragged tensor 存储；KV cache 用 BSR，B_r 对应 query 的 tile 大小，B_c 由 KV cache 管理算法决定（page 大小，可以小到 1，即 vector sparse）。kernel 支持任意的 (B_r, B_c)。

源码：分页格式的结构体 flashinferpaged_kv_t；通用的块稀疏接口 flashinferBlockSparseAttentionWrapper（它的 flashinferplan 直接接收 BSR 的 indptr 和 indices）、行块大小可变的 flashinferVariableBlockSparseAttentionWrapper。

可组合格式：共享前缀用大块

单一格式有个两难：B_r 越大，同一个块里的 query 越多，它们读进 shared memory 的 KV 可以共享；但 B_r 太大又会浪费（块里有很多空行）。

可组合格式（composable formats）的思路：同一个稀疏矩阵，拆成几个子矩阵，每个用最合适的块大小。最典型的是共享前缀：4 个请求共用同一段系统 prompt（页 0、1），后面各自有不同的后缀。

列是物理页。页 0、1 是共享前缀。左边：前缀被读了 4 遍，而且每个请求在自己的 CTA 里读，只能走全局内存或 L2。右边：前缀子矩阵是一个 4 × 2 的稠密块，4 个 query 在同一个 CTA 里，前缀只读进 shared memory 一次，4 个 query 共用。两部分各算出一个 attention state，再用 ⊕ 合并。

关键是：这不需要移动 KV cache 里的任何数据，只需要给两个子矩阵各算一套 indptr / indices。论文的实验（MLC-Engine，并行生成 n = 4）里，可组合格式让 ITL 降了 13.7%（8B）和 17.4%（70B）。

源码：flashinferMultiLevelCascadeAttentionWrapper（多层共享前缀），F2 讲过它最后用 merge_state 合并；讲解文章 Cascade Inference。

⑥ 设计 2：稀疏地读，稠密地算

块稀疏带来一个问题：Tensor Core 要求输入是规整的稠密小矩阵（至少 16 行），而 B_c 可能是 1（每个 KV 位置单独分散存放）。

FlashInfer 的做法分两步：先把分散的 KV 行“搬”成 shared memory 里连续的一块，再用稠密的 Tensor Core 计算。

• 地址怎么算：稀疏 KV 的地址来自 BSR 的 indices 数组；稠密 KV 的地址是行号的仿射变换。这和你在 L2 里写的 phys = block_tables[b][p]、kv_off = phys * stride + … 是同一件事。
• TMA 用不上：Hopper 的 TMA 只支持固定 stride 的访问，而稀疏 gather 需要任意的行地址。所以 FlashInfer 只在 KV 连续时用 TMA，其他情况退回 Ampere 风格的异步拷贝。
• 代价：论文附录 B 测了 page size = 1（最极端的稀疏）对比连续 KV：decode 几乎没有差别（1% 以内），prefill 约慢 10%（FA3 模板上稀疏加载占更多寄存器，只能用更小的 KV tile）。

源码：Hopper 上的稀疏加载主循环 flashinfersparse_mainloop.cuh（对比稠密的 flashinfermainloop.cuh）；FA2 模板在 flashinferprefill.cuh。

⑦ 设计 3：多种 tile 尺寸，以及 GQA 的 head-group 融合

按算术强度选 tile

标准的 FlashAttention-2 只有少数几种 tile（比如 128 × 64），对 A100 上的 prefill 最优，对 decode 却很浪费：decode 只有 1 个 query，128 行的 tile 里 127 行是空的。

FlashInfer 准备了 query 方向 {1, 16, 32, 64, 128} × KV 方向 {32, 64, 128} 的多种 tile，按两条规则选：

1. 算出这一批的平均 query 长度，选不小于它的最小 query tile。
2. 在寄存器和 shared memory 的限制下，选能让 SM 占用率最高的 KV tile。

query tile 为 1 时用 CUDA core（Tensor Core 至少要 16 行），其余用 Tensor Core；Hopper 的 FA3 模板用 64 的倍数（wgmma 的要求）。

源码里的选择函数是 flashinferFA2DetermineCtaTileQ：平均长度超过 64 选 128，16 到 64 之间选 64，再短选 16。调用处是 flashinferPrefillSplitQOKVIndptr 里按 avg_packed_qo_len 选。

head-group 融合：让 decode 也能用 Tensor Core

“平均 query 长度”用的是打包后的长度：flashinferpacked_qo_len = qo_len × gqa_group_size。这就是论文附录 A 的 head-group 融合，也就是你在 L2 里做的事：同一个 KV head 下的几个 query head，和 query 长度这一维合并成一个大的“行”维度。

GQA 组大小 g = 4，每个请求 2 个 query（比如 speculative decoding）：

2 个 token × 4 个 head = 8 行，全部读同一个 KV head：KV 只读进 shared memory 一次，8 行一起用。packed_qo_len = 2 × 4 = 8，所以选 16 行的 query tile，可以用 Tensor Core。query 足够长时（prefill），query 维度本身就够填满 tile，融合就不那么关键了。

⑧ 设计 4：可定制的 attention 模板 + JIT

观察：变体越来越多，但骨架都一样

GQA、MLA、滑动窗口、logits soft-cap（Gemma、Grok）、attention sink、自定义 mask、把 RoPE 融合进 attention、甚至不用 softmax 的 FlashSigmoid……为每一种手写一个专用 kernel 是不可持续的。但它们的骨架都是 FlashAttention：只是在几个固定的位置做了一点小修改。论文把这个设计空间写成：

f_epilogue( scan( f_logits( f_q(Q) · f_k(K) ) ) · f_v(V) )

做法：变体 = 几个固定位置的函数（functor）

受 FlexAttention 启发，FlashInfer 的模板在这些位置留了“插槽”：

用户写一段 CUDA 代码定义一个变体类，JIT 编译器把它填进模板，生成专用 kernel，编译好的 .so 缓存起来（F5 讲过 JIT 的缓存机制）。

一个真实的例子：attention sink（gpt-oss 等模型用到），完整定义在 flashinfervariants.py，核心是三段：

struct AttentionSink : AttentionVariantBase {
  REGISTER_LOGITS_MASK(params, batch_idx, qo_idx, kv_idx, qo_head_idx, kv_head_idx, {
    return (kv_idx + qo_len + window_left >= kv_len + qo_idx);   // 滑动窗口 + 因果
  })
  REGISTER_M_D_UPDATE(params, kv_tile_idx, qo_head_idx, m, d, scale, {
    // 第一个 KV 块时，把每个 head 的 sink 值当作一个额外的 logit 加进 online softmax 的 (m, d)
  })
  REGISTER_OUTPUT_TRANSFORM(params, output, batch_idx, qo_idx, qo_head_idx, m, d, scale, {
    return output * scale * d_rcp;                                // 用加了 sink 的分母归一化
  })
};

测试里调用它的方式（flashinfertest_attention_sink.py）：把变体名、变体代码、额外的张量 sink 和标量 sm_scale 作为 jit_args 传给 BatchPrefillWithRaggedKVCacheWrapper，就得到一个专用的 kernel。

默认的变体也是用同样的机制写的：flashinferDefaultAttention，flashinfersoft-cap 在 LogitsTransform 里做 tanh，flashinfer滑动窗口在 LogitsMask 里。注册宏在 flashinfervariant_helper.cuh。

为什么用 CUDA / CUTLASS 而不是 Triton（论文附录 C）：需要 warp specialization、TMA 这类 Triton 当时还不支持或处于实验阶段的硬件特性，以及寄存器级别的精细控制和直接写 PTX。论文附录 G 和 FlexAttention 的对比里，FlashInfer 在 H100 上的各种变体都更快，长序列时差距更大。

⑨ 设计 5：负载均衡调度、plan / run 与 CUDA graph

观察：一个请求一个 CTA 会严重不均衡

serving 时同一个 batch 里的请求长度差别很大。最朴素的做法是每个请求（或每个 query tile）一个 CTA，最长的那个决定整体时间，其他 SM 早早就闲着了。

做法：先切块，再用最小堆分配（论文算法 1）

1. 每个工作单元的代价：cost = α · l_q + β · l_kv。源码里是 flashinfercost_function = 2 · qo_len + kv_len。
2. 每块 KV 的最大长度 L_kv = 总工作量 ÷ CTA 数。比这个长的请求被切成几块。
3. 把所有块按长度从大到小排序。
4. 用一个最小堆维护每个 CTA 目前的总代价，每次把下一块分给最闲的 CTA。源码：flashinferMinHeap cta_cost_heap，pop 最小，加上 cost，再 push 回去。
5. 被切开的请求，各块的结果是 attention state，最后用 ⊕ 合并。

算一个例子：4 个 decode 请求，KV 长度 [1000, 100, 100, 50]，4 个 CTA，代价只看 l_kv。总量 1250，L_kv = ⌈1250 / 4⌉ = 313。请求 A 被切成 313、313、313、61 四块；排序后是 313、313、313、100、100、61、50。

分配过程：前三块 313 分给 CTA 0、1、2；之后 CTA 3 一直是最闲的，依次拿到 B（100）、C（100）、A₄（61）、D（50），总共 311。A 的四块各产生一个 attention state，最后用 ⊕ 合并成 A 的输出。

和 Stream-K 的区别：思路来自 GEMM 里的 Stream-K，但 Stream-K 用原子加合并部分结果，顺序不确定。serving 要求输出确定，所以 FlashInfer 不用原子操作，同样的长度输入总是产生同样的合并顺序（C2 讲过浮点原子加为什么不确定）。另外，不需要切的短请求直接写最终输出，不经过 workspace（论文附录 D.2）。

plan / run：每一步排一次班，所有层共用

调度要读每个请求的长度、排序、用堆分配，这些都在 CPU 上做（plan）；结果异步拷进 GPU 的 workspace；GPU 上的 kernel（run）只按计划执行。

怎么兼容 CUDA graph

CUDA graph 一旦捕获，kernel 的 grid 大小和参数（包括指针）就固定了。但每一步的调度结果都不一样。FlashInfer 的解决办法：

• persistent kernel：grid 大小固定（每个 SM 1–2 个 CTA），每个 CTA 循环处理 plan 分给它的工作列表，工作量变了，grid 不用变。源码：flashinferBlockBatchPagedAttentionPersistent（attention）和 flashinferBlockBatchReductionPersistent（⊕ 合并），每个 CTA 用 flashinferget_block_coord(params, work_idx) 从 plan 里取出自己的下一块工作。
• workspace 的每一段都在固定的偏移：按上限一次分配好，所以传给 kernel 的指针每一步都一样，只是指向的内容变了。
• plan 本身在 CPU 上，不进 graph；graph 里只有 run。

源码：Python 端 flashinferBatchDecodeWithPagedKVCacheWrapper.plan（F1 用过）、flashinferBatchPrefillWithPagedKVCacheWrapper；C++ 端的计划 flashinferPrefillPlanInfo、flashinferHolisticPlanInfo 和 flashinferTwoStageHolisticPlan。

⑩ 编程接口：serving 框架怎么用它

论文清单 1 的流程（简化）：

workspace = torch.empty(...)                          # 预先分配，地址固定
attn = AttentionWrapper(attn_spec, task_info, workspace)   # 初始化时 JIT 编译 kernel
# 捕获阶段：每种配置先做一次 plan，再把所有层的 run 捕获进 CUDA graph
attn.plan(seqlen_info)
with torch.cuda.graph(g):
    for layer in layers:
        ...; attn.run(...); ...
# 生成循环：每一步只做 plan + replay
while not finished:
    seqlen_info.update()
    attn.plan(seqlen_info)      # CPU 上排班，写进 workspace
    g.replay()                  # GPU 上所有层按新计划执行

不同平均 query 长度、不同可组合格式的配置会编译成不同的 kernel、捕获成不同的 graph，运行时由 serving 框架按当前情况挑一个。论文里列出的集成方是 vLLM、SGLang 和 MLC-Engine。vLLM 的调用见 vllmflashinfer.py 里的 plan，SGLang 的见 sglangflashinfer_backend.py。

⑪ 效果，以及和这套题的对应

论文（FlashInfer v0.2，A100 / H100）报告的结果：

你在这套题里已经亲手写过的部分：

⑫ 局限，以及之后的发展

论文自己指出的局限：

• 只做 forward。训练需要可定制的 backward 模板，论文把它列为未来工作。
• 稀疏 KV 在 Hopper 上用不了 TMA，只能走异步拷贝，prefill 约慢 10%。B_c 很大（比如 128）时可以用 TMA，但会牺牲格式的灵活性。

论文发表之后，FlashInfer 的范围扩展了很多（以本仓库钉住的版本为准）：Blackwell 上的 attention 后端（TensorRT-LLM、cuDNN、CuTe DSL）、MLA、采样（F4）、norm 和激活（F3）、FP8 / FP4 的 GEMM 和 MoE。核心的设计原则没有变：统一的数据格式 + 编译期定制 + 运行期调度。项目自己的说明见 flashinferCLAUDE.md。

源码与资料索引

论文与文章：FlashInfer 论文（arXiv 2501.01005） · Cascade Inference 博客 · Sorting-Free Sampling 博客 · FlashInfer 文档 · FlexAttention

数据结构：flashinferstate_t（attention state） · flashinferpaged_kv_t · flashinferBlockSparseAttentionWrapper · flashinferMultiLevelCascadeAttentionWrapper

定制：flashinferDefaultAttention · flashinferREGISTER_* 宏 · flashinferAttentionSink 变体

调度：flashinfercost_function · flashinfer最小堆分配 · flashinferFA2DetermineCtaTileQ · flashinferpersistent kernel

相关讲解：T5 · L2 · L3 · F1 · F2 · F5