手写 mini-sglang 阶段四 · 更快、更大

GPU 注意力后端：FlashInfer 与 FlashAttention

参考后端用 Python 循环逐个请求算注意力，语义清楚但慢。GPU 上，mini-sglang 用两个专门的库：FlashInfer（decode 很强，prefill 也好用）和 FlashAttention 3（Hopper 上 prefill 最快）。它们的接口不同——FlashInfer 要先 plan 再 run，FlashAttention 直接吃页号表——但都能表达第 5 章那些元数据。这一章按官方的写法实现两个后端，并用同接口的 PyTorch 假实现在 CPU 上验证参数语义。

自测：能答上来就可以跳过本章

1. FlashInfer 的 plan 和 run 分别做什么？为什么 plan 每个 batch 只做一次？
2. mini-sglang 用 FlashInfer 时把 KV 池当成 page size = 1，indices 里存的是什么？last_page_len 为什么恒为 1？
3. FlashAttention 的 page_table 里存的是 token 位置还是页号？怎样从 mini-sglang 的全局 page table 得到它？
4. 在 CPU 上怎样验证一个只能在 GPU 上运行的后端？

自测参考答案（先自己答，再展开对照）

1. plan 根据这个 batch 的元数据（每个请求的长度、页表）做调度规划：划分工作、分配临时缓冲区、把元数据拷到 GPU；run 在每一层按规划执行计算。同一个 batch 的所有层元数据相同，所以只 plan 一次（推迟到第一层时才做）。
2. indices 里存的是每个 token 的 KV 在池中的位置（page size 为 1 时页号就是 token 位置），也就是 page table 的一行；每一页只有一个 token，最后一页总是满的，last_page_len 恒为 1。
3. 页号。从全局 page table（按 token 存位置）每隔 page size 取一个位置，再除以 page size，就得到每一页的页号。
4. 写一个接口与真实库相同的假实现：把两个库的参数语义翻译成同一个参考实现（Python + SDPA），在 CPU 上跑后端代码，结果要和参考后端完全一致。

本章要写的文件：attention/fi.py、attention/fa.py，以及测试用的 tests/fakes/flashinfer/、tests/fakes/sgl_kernel/。

到这一章为止的文件：75 个（新建 2，修改 0，提前引入 0）

minisgl/
├── __init__.py
├── core.py
├── env.py
├── __main__.py
├── shell.py
├── attention/
│   ├── __init__.py
│   ├── base.py
│   ├── utils.py
│   ├── torch_backend.py
│   ├── fi.py  ← 本章新建
│   ├── fa.py  ← 本章新建
├── distributed/  （3 个文件，不变）
├── engine/  （5 个文件，不变）
├── kernel/  （2 个文件，不变）
├── kvcache/  （5 个文件，不变）
├── layers/  （8 个文件，不变）
├── llm/  （2 个文件，不变）
├── message/  （5 个文件，不变）
├── models/  （7 个文件，不变）
├── moe/  （2 个文件，不变）
├── scheduler/  （9 个文件，不变）
├── server/  （4 个文件，不变）
├── tokenizer/  （4 个文件，不变）
├── utils/  （8 个文件，不变）

这一步的 main：examples/ch17_gpu_attention.py——它只用到上面这些文件；python tools/steps.py check 会逐章搭出这棵树、跑这个 main。

FlashInfer：plan 与 run

FlashInfer 为分页 KV 缓存提供了"批量 prefill"和"批量 decode"两种 wrapper。用法分两步：

• plan(...)：传入这个 batch 的形状信息——每个请求的 query 范围（qo_indptr）、每个请求的 KV 由哪些页组成（kv_indptr 划分 kv_indices）、最后一页用了多少（last_page_len）。FlashInfer 在 CPU 上据此算出工作划分（哪个线程块处理哪段 KV），异步拷到 GPU；
• run(q, paged_kv_cache)：每层调用一次，按 plan 好的方案计算。

28 层的形状完全相同，plan 只需要一次。mini-sglang 把 plan 推迟到第一层的 forward 里（_initialize_metadata_once），而不是在 prepare_metadata 里：

minisgl/attention/fi.py · FlashInferBackend.forward

def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, layer_id: int,
            batch: Batch) -> torch.Tensor:
    metadata = batch.attn_metadata
    assert isinstance(metadata, FIMetadata)
    self._initialize_metadata_once(metadata)  # 第一层时 plan，之后各层直接 run
    self.kvcache.store_kv(k, v, batch.out_loc, layer_id)

    def flatten(cache: torch.Tensor) -> torch.Tensor:  # [页, 页大小, H, D] -> [token, 1, H, D]
        return cache.view(-1, 1, cache.shape[2], cache.shape[3])

    kv = (flatten(self.kvcache.k_cache(layer_id)), flatten(self.kvcache.v_cache(layer_id)))
    return metadata.wrapper.run(q=q, paged_kv_cache=kv)

minisgl/attention/fi.py · FlashInferBackend.prepare_metadata

def prepare_metadata(self, batch: Batch) -> None:
    reqs = batch.padded_reqs
    bs = len(reqs)
    seqlens_q = [r.extend_len for r in reqs]
    seqlens_k = [r.device_len for r in reqs]
    kw = dict(dtype=torch.int32, pin_memory=pin(self.device))
    seq_lens_cpu = torch.tensor(seqlens_k, **kw)
    cu_k = torch.tensor([0] + seqlens_k, **kw).cumsum_(dim=0)
    if max(seqlens_q) == 1:  # decode：每个请求 1 个 query
        cu_q = torch.arange(0, bs + 1, **kw)
    elif all(r.cached_len == 0 for r in reqs):  # 没有命中缓存的 prefill：q 与 k 等长
        cu_q = cu_k
    else:
        cu_q = torch.tensor([0] + seqlens_q, **kw).cumsum_(dim=0)
    page_table = get_global_ctx().page_table
    batch.attn_metadata = FIMetadata(
        cu_seqlens_q_cpu=cu_q,
        cu_seqlens_k_cpu=cu_k,
        cu_seqlens_q_gpu=cu_q.to(self.device, non_blocking=True),
        indices=torch.cat([page_table[r.table_idx, : r.device_len] for r in reqs]),
        last_page_len_cpu=self._get_ones_cpu(bs),
        num_qo_heads=self.qo_head_local,
        num_kv_heads=self.kv_head_local,
        head_dim=self.config.head_dim,
        seq_lens_cpu=seq_lens_cpu,
        dtype=self.kvcache.dtype,
        wrapper=self.decode_wrapper if batch.is_decode else self.prefill_wrapper,
    )

几个值得注意的设计：

• 把 KV 池当成 page size = 1：forward 里把 KV 池 view 成 [token 数, 1, 头, 维]，于是全局 page table 的一行就是 kv_indices，last_page_len 恒为 1。这样无论系统的 page size 是多少，FlashInfer 后端都不用做任何换算（第 4 章"page table 按 token 存位置"的好处）；
• 两个快路径：纯 decode 时 qo_indptr = [0, 1, …, bs]；没有前缀命中的 prefill 时每个请求的 query 数等于 KV 长度，qo_indptr 直接复用 kv_indptr；
• 锁页内存与事件：plan 会复用一块锁页内存做异步拷贝，所以在下一次 plan 修改它之前，要先等上一次拷贝完成（last_event.synchronize()）；
• Tensor Core：GQA 组大于等于 4 时，decode 实际上是一个个小矩阵乘，用 Tensor Core 版本的 decode kernel 更快。

minisgl/attention/fi.py · FlashInferBackend._initialize_metadata_once

def _initialize_metadata_once(self, metadata: FIMetadata) -> None:
    if metadata.initialized:
        return
    from flashinfer import BatchDecodeWithPagedKVCacheWrapper

    metadata.initialized = True
    # plan 会复用一块锁页内存做异步拷贝，必须等上一次拷贝结束再改写它
    self.last_event.synchronize()
    if isinstance(metadata.wrapper, BatchDecodeWithPagedKVCacheWrapper):
        metadata.wrapper.plan(
            indptr=metadata.cu_seqlens_k_cpu, indices=metadata.indices,
            last_page_len=metadata.last_page_len_cpu, num_qo_heads=metadata.num_qo_heads,
            num_kv_heads=metadata.num_kv_heads, head_dim=metadata.head_dim, page_size=1,
            pos_encoding_mode="NONE", seq_lens=metadata.seq_lens_cpu,
            data_type=metadata.dtype, q_data_type=metadata.dtype,
            kv_data_type=metadata.dtype, non_blocking=True,
        )
    else:
        metadata.wrapper.plan(
            qo_indptr=metadata.cu_seqlens_q_cpu, paged_kv_indptr=metadata.cu_seqlens_k_cpu,
            paged_kv_indices=metadata.indices, paged_kv_last_page_len=metadata.last_page_len_cpu,
            num_qo_heads=metadata.num_qo_heads, num_kv_heads=metadata.num_kv_heads,
            head_dim_qk=metadata.head_dim, page_size=1, pos_encoding_mode="NONE",
            seq_lens=metadata.seq_lens_cpu, q_data_type=metadata.dtype,
            kv_data_type=metadata.dtype, non_blocking=True, causal=True,
        )
    self.last_event.record()

FlashAttention：页号表

sgl_kernel 里的 flash_attn_with_kvcache（FlashAttention 3）直接读分页的 KV 池，需要：page_table（每个请求的页号）、cache_seqlens（每个请求的 KV 长度）、cu_seqlens_q（变长 query 的划分）。没有 plan 步骤。

minisgl/attention/fa.py · FlashAttentionBackend.prepare_metadata

def prepare_metadata(self, batch: Batch) -> None:
    reqs = batch.padded_reqs
    bs = len(reqs)
    seqlens_q = [r.extend_len for r in reqs]
    seqlens_k = [r.device_len for r in reqs]
    max_seqlen_k, max_seqlen_q = max(seqlens_k), max(seqlens_q)
    kw = dict(dtype=torch.int32, pin_memory=pin(self.device))
    cache_seqlens = torch.tensor(seqlens_k, **kw).to(self.device, non_blocking=True)
    cu_k = torch.tensor([0] + seqlens_k, **kw).cumsum_(dim=0).to(self.device, non_blocking=True)
    if max_seqlen_q == 1:
        cu_q = torch.arange(0, bs + 1, dtype=torch.int32, device=self.device)
    elif all(r.cached_len == 0 for r in reqs):
        cu_q = cu_k
    else:
        cu_q = torch.tensor([0] + seqlens_q, **kw).cumsum_(dim=0).to(self.device,
                                                                      non_blocking=True)
    # 全局 page table 按 token 存位置；每隔 page_size 取一个，再除以 page_size 就是页号
    page_table = get_global_ctx().page_table
    table = torch.stack([page_table[r.table_idx, :max_seqlen_k : self.page_size] for r in reqs])
    if self.page_size > 1:
        table = table.div(self.page_size, rounding_mode="floor")
    batch.attn_metadata = FAMetadata(cu_k, cu_q, cache_seqlens, max_seqlen_k, max_seqlen_q, table)

全局 page table 按 token 存位置，每隔 page size 取一个位置就是每页的起点，再除以 page size 就是页号：page_table[r, :max_len:page_size] // page_size。

在 CPU 上验证

这两个库都只能在 NVIDIA GPU 上运行。为了验证后端代码传给它们的参数是对的，我们为每个库写了一个"同接口的假实现"，放在 tests/fakes/ 下，测试时加到 import 路径最前面：

tests/fakes/flashinfer/__init__.py · BatchPrefillWithPagedKVCacheWrapper

class BatchPrefillWithPagedKVCacheWrapper(_Base):
    def plan(self, qo_indptr, paged_kv_indptr, paged_kv_indices, paged_kv_last_page_len,
             num_qo_heads, num_kv_heads, head_dim_qk, page_size, causal=False, **kwargs: Any):
        assert causal
        self._plan(qo_indptr, paged_kv_indptr, paged_kv_indices, paged_kv_last_page_len, head_dim_qk)

tests/fakes/sgl_kernel/flash_attn.py · flash_attn_with_kvcache

def flash_attn_with_kvcache(q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor,
                            page_table: torch.Tensor, cache_seqlens: torch.Tensor,
                            cu_seqlens_q: torch.Tensor, cu_seqlens_k_new: torch.Tensor | None = None,
                            max_seqlen_q: int = 1, softmax_scale: float | None = None,
                            causal: bool = True, **kwargs) -> torch.Tensor:
    assert causal
    page_size = k_cache.shape[1]
    lens = cache_seqlens.tolist()
    pages = [page_table[i, : (n + page_size - 1) // page_size] for i, n in enumerate(lens)]
    scale = softmax_scale if softmax_scale is not None else q.shape[-1] ** -0.5
    return paged_attention(q, k_cache, v_cache, cu_seqlens_q.tolist(), pages, lens, scale)

两者都把参数翻译成"每个请求的 query 范围、KV 页列表、KV 长度"，交给同一个参考实现：

tests/fakes/_fake_ref.py · paged_attention

def paged_attention(q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor,
                    q_starts: List[int], kv_pages: List[torch.Tensor], kv_lens: List[int],
                    scale: float) -> torch.Tensor:
    """q: [总 query 数, Hq, D]；k_cache: [页数, 页大小, Hkv, D]。

    第 i 个请求的 query 是 q[q_starts[i]:q_starts[i+1]]，KV 是 kv_pages[i] 这些页拼起来的前 kv_lens[i]
    个 token。query 是序列的最后 n 个位置（因果掩码右下角对齐）。
    """
    out = torch.empty_like(q)
    for i, (pages, kv_len) in enumerate(zip(kv_pages, kv_lens)):
        qs, qe = q_starts[i], q_starts[i + 1]
        n = qe - qs
        if n == 0:
            continue
        k = k_cache[pages.long()].flatten(0, 1)[:kv_len].transpose(0, 1)  # [Hkv, L, D]
        v = v_cache[pages.long()].flatten(0, 1)[:kv_len].transpose(0, 1)
        q_pos = torch.arange(kv_len - n, kv_len, device=q.device).unsqueeze(1)
        mask = torch.arange(kv_len, device=q.device).unsqueeze(0) <= q_pos
        o = F.scaled_dot_product_attention(q[qs:qe].transpose(0, 1), k, v, attn_mask=mask,
                                           scale=scale, enable_gqa=True)
        out[qs:qe] = o.transpose(0, 1)
    return out

这样，后端代码里任何一个参数传错（比如 indptr 少了开头的 0、页号没有除以 page size、decode 时 qo_indptr 写错），假实现算出的结果就会与参考后端不同，测试失败。这套假实现本身也是理解 FlashInfer 和 FlashAttention 接口语义的最短路径。

examples/ch17_gpu_attention.py

"""第 17 章：FlashInfer 与 FlashAttention 后端拿到的元数据长什么样。

CPU 上没有这两个库，这里把 tests/fakes（同接口的 PyTorch 假实现）放到 import 路径上，
官方写法的后端代码原样运行，再与参考后端比较。
"""

import sys

sys.path.insert(0, "tests/fakes")

import torch  # noqa: E402
from minisgl.core import Batch, Req, SamplingParams, reset_global_ctx  # noqa: E402
from minisgl.distributed import DistributedInfo, reset_tp_info  # noqa: E402
from minisgl.engine import Engine, EngineConfig  # noqa: E402

PROMPTS = [[785, 6722, 315, 9625, 374], [16, 488, 220, 16, 284, 220]]


def build(backend: str, page_size: int) -> Engine:
    reset_global_ctx()
    reset_tp_info()
    e = Engine(EngineConfig(model_path="models/Qwen3-0.6B", tp_info=DistributedInfo(0, 1),
                            dtype=torch.float32, attention_backend=backend, page_size=page_size,
                            max_running_req=2, num_page_override=64 // page_size, max_seq_len_override=32))
    e.page_table[0, :16] = torch.arange(16, 32)  # 请求 0：第 16～31 个位置
    e.page_table[1, :16] = torch.arange(40, 56)  # 请求 1：第 40～55 个位置
    return e


def batch_of(engine: Engine, phase: str, reqs):
    b = Batch(reqs=reqs, phase=phase)
    b.padded_reqs = reqs
    b.positions = torch.cat([torch.arange(r.cached_len, r.device_len) for r in reqs]).int()
    b.input_ids = torch.cat([r.input_ids[r.cached_len:r.device_len] for r in reqs])
    b.out_loc = torch.cat([engine.page_table[r.table_idx, r.cached_len:r.device_len] for r in reqs])
    engine.attn_backend.prepare_metadata(b)
    return b


def reqs(cached):
    return [Req(input_ids=torch.tensor(p, dtype=torch.int32), table_idx=i, cached_len=c, output_len=4,
                uid=i, sampling_params=SamplingParams(), cache_handle=None)
            for i, (p, c) in enumerate(zip(PROMPTS, cached))]


results = {}
for backend, ps in (("torch", 1), ("fi", 1), ("fa", 4)):
    engine = build(backend, ps)
    with engine.ctx.forward_batch(b := batch_of(engine, "prefill", reqs([0, 0]))):
        engine.model.forward()  # 先把两个请求的 KV 都算好
    rs = reqs([3, 5])  # 请求 0 命中前 3 个 token，请求 1 只剩最后 1 个（像 decode）
    b = batch_of(engine, "prefill", rs)
    m = b.attn_metadata
    if backend == "fi":
        print("FlashInfer（按 page_size=1 使用 KV 池）:")
        print("  qo_indptr      =", m.cu_seqlens_q_cpu.tolist())
        print("  kv_indptr      =", m.cu_seqlens_k_cpu.tolist())
        print("  kv_indices     =", m.indices.tolist())
        print("  last_page_len  =", m.last_page_len_cpu.tolist())
    if backend == "fa":
        print("FlashAttention（page_size=4，page table 存页号）:")
        print("  cu_seqlens_q   =", m.cu_seqlens_q.tolist())
        print("  cache_seqlens  =", m.cache_seqlens.tolist())
        print("  page_table     =", m.page_table.tolist())
    with engine.ctx.forward_batch(b):
        results[backend] = engine.model.forward()
    if backend == "fi":
        print("  prefill wrapper 的 plan 次数:", engine.attn_backend.prefill_wrapper.plan_count,
              "（两个 batch 各一次；每个 batch 的 28 层共用一次 plan）")
    engine.shutdown()
for name in ("fi", "fa"):
    print(f"{name} 与参考后端 logits 的最大误差: {(results[name] - results['torch']).abs().max().item():.2e}")

运行结果

FlashInfer（按 page_size=1 使用 KV 池）:
  qo_indptr      = [0, 2, 3]
  kv_indptr      = [0, 5, 11]
  kv_indices     = [16, 17, 18, 19, 20, 40, 41, 42, 43, 44, 45]
  last_page_len  = [1, 1]
  prefill wrapper 的 plan 次数: 2 （两个 batch 各一次；每个 batch 的 28 层共用一次 plan）
FlashAttention（page_size=4，page table 存页号）:
  cu_seqlens_q   = [0, 2, 3]
  cache_seqlens  = [5, 6]
  page_table     = [[4, 5], [10, 11]]
fi 与参考后端 logits 的最大误差: 0.00e+00
fa 与参考后端 logits 的最大误差: 0.00e+00

两个请求：请求 0 命中前 3 个 token、本轮算 2 个；请求 1 前 5 个在缓存里、本轮算 1 个。FlashInfer 拿到的是 token 级的 kv_indices（16～20 和 40～45），last_page_len 全为 1；FlashAttention 在 page size 为 4 时拿到的是页号（位置 16 在第 4 页，位置 40 在第 10 页）。两个后端的 logits 与参考后端完全相同。

官方实现

• attention/fi.py · FlashInferBackend（第 80～271 行）
• attention/fa.py · FlashAttentionBackend（第 36～136 行）
• 官方还有 TensorRT-LLM 的后端 attention/trtllm.py（Blackwell 上默认使用，只支持 page size 16、32、64），我们没有实现。

我们的两个后端与官方的逻辑一致，只把 pin_memory=True 和 torch.cuda.Event 换成了设备抽象里的函数。

测试

tests/test_ch17_gpu_backends.py · test_backend_logits_match_torch_backend

@pytest.mark.parametrize("backend", ["fi", "fa"])
def test_backend_logits_match_torch_backend(fake_kernels, backend):
    """同一个 batch（一个命中前缀的 prefill + 一个 decode），FlashInfer/FlashAttention 与参考后端一致。"""
    from transformers import AutoTokenizer

    tok = AutoTokenizer.from_pretrained(QWEN3)
    ids = [tok(p).input_ids for p in PROMPTS[:2]]
    logits = {}
    for name in ("torch", backend):
        engine = build_engine(QWEN3, attention_backend=name)
        identity_page_table(engine, 2)
        forward(engine, [new_req(ids[0][:3], 0), new_req(ids[1], 1)], "prefill")  # 先把 KV 填好
        a = new_req(ids[0], 0, cached_len=3)  # 命中前 3 个 token，再算剩下的
        b = new_req(ids[1] + [100], 1, cached_len=len(ids[1]))  # decode 一步
        logits[name] = forward(engine, [a, b], "prefill")
        engine.shutdown()
    assert (logits["torch"] - logits[backend]).abs().max().item() < 1e-4

tests/test_ch17_gpu_backends.py 还检查了"每个 batch 只 plan 一次"，并用 FlashInfer（page size 1）、FlashAttention（page size 4）、fa,fi 组合三种配置，在重叠调度和分块 prefill 下端到端地与 Hugging Face 对比。

怎么讲清楚

讲 GPU 注意力后端：FlashInfer 分 plan 和 run——plan 根据这个 batch 的长度分布算好任务划分和临时缓冲区（CPU 上的工作，每个 batch 一次，推迟到第一层执行），run 每层调用一次；mini-sglang 把 KV 池当成 page size 1，page table 的一行就是 kv_indices，last_page_len 恒为 1。FlashAttention 直接吃页号表、KV 长度和 cu_seqlens_q，页号从全局 page table 里每隔"页大小"取一个位置再除以页大小。没有 GPU 时，用同接口的假实现把两个库的参数语义翻译成参考实现，验证后端的元数据构造是否正确。

练习

1. 如果要让 FlashInfer 使用真正的 page size（例如 16）而不是展平成 1，prepare_metadata 要怎么改？last_page_len 怎么算？
2. 在假 FlashInfer 里故意把 qo_indptr 的解释改错（比如忽略开头的 0），哪个测试会失败？
3. 为什么 CUDA Graph 只和 decode 后端打交道，而 FlashAttention 的 prefill 不需要 CUDA Graph？

参考答案

1. kv_indices 改为页号：page_table[r, :device_len:ps] // ps；kv_indptr 按页数累加（ceil(device_len / ps)）；last_page_len = device_len - (页数 - 1) * ps；forward 里不再 view 成 page size 1，而是直接传 [页数, ps, 头, 维] 的 KV 池。
2. test_backend_logits_match_torch_backend[fi] 和所有用 fi 的端到端测试都会失败——这正是假实现的价值。
3. prefill 的 batch 形状（token 数、各请求长度）千变万化，无法为每种形状录制 graph；prefill 的计算量大，kernel 启动开销占比小，也不需要。decode 的形状只由批大小决定，且每步计算很少，启动开销占比大。

小结

• FlashInfer：plan（每个 batch 一次，推迟到第一层）+ run（每层）；mini-sglang 把 KV 池当成 page size 1，page table 的一行就是 kv_indices。
• FlashAttention：直接吃页号表、KV 长度和 cu_seqlens_q，页号从全局 page table 每隔 page size 取一个再除以 page size。
• 同接口的假实现把两个库的参数语义翻译成同一个参考实现，在 CPU 上验证后端代码，结果与参考后端完全一致。

本章练习题（1 道）

在浏览器里直接写代码、跑测试；也可以在本地用 python practice/judge.py 判题。

• 200. FlashInfer 的分页元数据（中等）

发现错误或有更好的写法？在 GitHub 上编辑此页 · 提交问题