手写 mini-sglang 阶段一 · 算得对：把骨架换成正式结构

注意力后端与参考实现

注意力是引擎里唯一需要"知道 batch 结构"的计算：每个请求长度不同、KV 散落在池子的不同页里、有的请求前缀已经缓存。mini-sglang 把这些复杂性收进"注意力后端"：每个 batch 调用一次 prepare_metadata 把请求信息整理成 kernel 需要的格式，每层调用一次 forward 写 KV、算注意力。官方有 FlashInfer、FlashAttention、TensorRT-LLM 三个 GPU 后端；这一章先写一个 PyTorch 参考后端，把语义讲透，第 17 章再接 GPU 后端。

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

1. cu_seqlens_q = [0, 6, 9, 10] 表示什么？
2. 一个请求 KV 总长 7、本轮 3 个 query，因果掩码怎样对齐？第一个 query 能看到哪些 key？
3. 为什么元数据每个 batch 准备一次，而不是每层准备一次？
4. 写 KV 必须在算注意力之前还是之后？为什么？

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

1. 本轮 batch 里有 3 个请求，各自的 query 数是 6、3、1（前缀和的形式：第 i 个请求的 query 在 [cu[i], cu[i+1])）。
2. 按右下角对齐：3 个 query 是这个请求的最后 3 个位置（4、5、6）。第一个 query 在位置 4，能看到 key 0～4 共 5 个。
3. 同一个 batch 里，每一层的变长信息、KV 长度、页表都完全一样，准备一次所有层共用；每层重新准备是纯粹的浪费。
4. 之前：本轮新 token 自己的 K、V 也要参与注意力（因果注意力里每个 token 要看到自己），所以先写进 KV 池，再从池里读出全部 KV 计算。

本章要写的文件：attention/base.py、attention/utils.py、attention/torch_backend.py、attention/__init__.py。

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

minisgl/
├── __init__.py
├── core.py
├── attention/
│   ├── __init__.py  ← 本章新建
│   ├── base.py  ← 本章新建
│   ├── utils.py  ← 本章新建
│   ├── torch_backend.py  ← 本章新建
├── distributed/  （3 个文件，不变）
├── kernel/  （2 个文件，不变）
├── kvcache/  （4 个文件，不变）
├── layers/  （8 个文件，不变）
├── models/  （7 个文件，不变）
├── scheduler/  （2 个文件，不变）
├── utils/
│   ├── __init__.py
│   ├── misc.py
│   ├── hf.py
│   ├── torch_utils.py
│   ├── registry.py
│   ├── logger.py
│   ├── device.py  ← 提前引入（第 6 章讲）

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

接口

minisgl/attention/base.py · BaseAttnBackend

class BaseAttnBackend(ABC):
    @abstractmethod
    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, layer_id: int,
                batch: Batch) -> torch.Tensor:
        """先把本轮的 k、v 写进 KV 池，再计算注意力。q: [T, Hq, D]，返回 [T, Hq, D]。"""

    @abstractmethod
    def prepare_metadata(self, batch: Batch) -> None:
        """每个 batch 调用一次（而不是每层一次）：把请求信息整理成 kernel 需要的格式。"""

    # 以下三个方法服务于 CUDA Graph（第 18 章）
    @abstractmethod
    def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None: ...

    @abstractmethod
    def prepare_for_capture(self, batch: Batch) -> None: ...

    @abstractmethod
    def prepare_for_replay(self, batch: Batch) -> None: ...

- prepare_metadata(batch)：调度器在每个 batch 开始前调用一次，结果放在 batch.attn_metadata。模型有 28 层，每层的注意力形状完全相同，元数据只需准备一次；GPU 后端还会在这里把数据从 CPU 拷到 GPU、或者做 FlashInfer 的 plan。 - forward(q, k, v, layer_id, batch)：每层调用。先写后读：把本轮新 token 的 k、v 写进 KV 池，再从池子里读出每个请求的全部 KV 算注意力。顺序不能反，因为本轮的 query 也要看到本轮的 key（因果掩码的对角线）。 - 最后三个方法服务于 CUDA Graph，第 18 章再讲。

元数据还要提供 get_last_indices(bs)：prefill 时每个请求最后一个 token 在展平后的 q 里的下标，给 LM head 用（第 2 章）。

变长、分页、带前缀的注意力

batch 里的张量是各请求首尾相接的一维数组。用一个例子说明注意力需要哪些信息：请求 A 从头 prefill 6 个 token；请求 B 前 4 个 token 命中前缀缓存，本轮算 3 个；请求 C 已经有 8 个 token，本轮 decode 1 个。

• 每个请求本轮的 query 在哪：cu_seqlens_q = [0, 6, 9, 10]（cumulative sequence lengths），第 i 个请求是 q[cu[i]:cu[i+1]]；
• 每个请求的 KV 有多长：cache_seqlens = [6, 7, 9]，即各请求的 device_len；
• 每个请求的 KV 在池中哪里：page table 的对应行，取前 cache_seqlens[i] 个位置；

一个混合 batch 的 query 划分、KV 来源和因果掩码

• 因果掩码怎么对齐：本轮的 n 个 query 是序列的最后 n 个位置。KV 长 7、query 3 个时，第一个 query 的绝对位置是 4，能看到位置 0～4 的 key。这叫"右下角对齐"：把 [n, kv_len] 的掩码矩阵画出来，它的对角线落在右下角。

参考后端就是把这段描述直接翻译成代码：

minisgl/attention/torch_backend.py · TorchAttnBackend.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, TorchAttnMetadata)
    self.kvcache.store_kv(k, v, batch.out_loc, layer_id)  # 先写后读：本轮的 token 也要参与注意力
    k_pool = self.kvcache.k_cache(layer_id).flatten(0, 1)  # [num_tokens, Hkv, D]
    v_pool = self.kvcache.v_cache(layer_id).flatten(0, 1)
    cu_q = metadata.cu_seqlens_q.tolist()
    seq_lens = metadata.cache_seqlens.tolist()
    out = torch.empty_like(q)
    for i, kv_len in enumerate(seq_lens):
        qs, qe = cu_q[i], cu_q[i + 1]
        n = qe - qs  # 本轮这个请求有 n 个 query，它们是序列的最后 n 个位置
        if n == 0:
            continue
        loc = metadata.page_table[i, :kv_len].long()
        qi = q[qs:qe].transpose(0, 1)  # [Hq, n, D]
        ki = k_pool[loc].transpose(0, 1)  # [Hkv, kv_len, D]
        vi = v_pool[loc].transpose(0, 1)
        # 第 r 个 query 的绝对位置是 kv_len - n + r，它能看到位置不超过自己的所有 key
        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
        oi = F.scaled_dot_product_attention(qi, ki, vi, attn_mask=mask, scale=self.scale,
                                            enable_gqa=True)
        out[qs:qe] = oi.transpose(0, 1)
    return out

对每个请求：按 page table 从池中取出它的全部 K、V，构造右下角对齐的因果掩码，调用 scaled_dot_product_attention（enable_gqa=True 让 16 个 q 头共享 8 个 KV 头）。这是一个 Python 循环，速度不快，但完全没有歧义，CPU 上也能跑。

元数据的准备：

minisgl/attention/torch_backend.py · TorchAttnBackend.prepare_metadata

def prepare_metadata(self, batch: Batch) -> None:
    reqs = batch.padded_reqs
    seqlens_q = [req.extend_len for req in reqs]
    seqlens_k = [req.device_len for req in reqs]
    max_seqlen_k = max(seqlens_k)
    page_table = get_global_ctx().page_table
    use_pin = pin(self.device)
    batch.attn_metadata = TorchAttnMetadata(
        cu_seqlens_q=make_cu_seqlens(seqlens_q, self.device, use_pin),
        cache_seqlens=torch.tensor(seqlens_k, dtype=torch.int32, pin_memory=use_pin)
        .to(self.device, non_blocking=True),
        page_table=torch.stack([page_table[r.table_idx, :max_seqlen_k] for r in reqs]),
        max_seqlen_q=max(seqlens_q),
    )

page table 取每个请求所在行的前 max_seqlen_k 列，拼成 [bs, max_seqlen_k]，这与官方 FlashAttention 后端的做法相同。用的是 batch.padded_reqs 而不是 batch.reqs：CUDA Graph 补齐进来的 dummy 请求也要有元数据。

把上面的例子跑一遍，和"不分页、不变长的完整因果注意力"对比：

examples/ch05_attention.py

"""第 5 章：一个混合 batch 的注意力元数据，以及参考后端与"不分页的完整注意力"的对比。

请求 A：6 个 token 从头 prefill；请求 B：前 4 个 token 命中缓存，再算 3 个；请求 C：decode 一步。
"""

import torch
import torch.nn.functional as F
from minisgl.attention.torch_backend import TorchAttnBackend
from minisgl.core import Batch, Context, Req, SamplingParams, set_global_ctx
from minisgl.distributed import set_tp_info
from minisgl.kvcache.mha_pool import MHAKVCache
from minisgl.models import ModelConfig
from minisgl.utils import cached_load_hf_config

set_tp_info(0, 1)
H, HKV, D = 16, 8, 128
ctx = Context(page_size=1)
set_global_ctx(ctx)
ctx.kv_cache = MHAKVCache(HKV, num_layers=1, head_dim=D, num_pages=64, page_size=1,
                          dtype=torch.float32, device=torch.device("cpu"))
ctx.page_table = torch.zeros((3, 16), dtype=torch.int32)
ctx.page_table[0, :6] = torch.arange(0, 6)
ctx.page_table[1, :7] = torch.tensor([40, 41, 42, 43, 20, 21, 22])  # 前 4 个位置来自前缀缓存
ctx.page_table[2, :9] = torch.arange(50, 59)
backend = TorchAttnBackend(ModelConfig.from_hf(cached_load_hf_config("models/Qwen3-0.6B")))


def req(n, cached, row):
    return Req(input_ids=torch.zeros(n, dtype=torch.int32), table_idx=row, cached_len=cached,
               output_len=4, uid=row, sampling_params=SamplingParams(), cache_handle=None)


torch.manual_seed(0)
lens, cached = [6, 7, 9], [0, 4, 8]
full = [(torch.randn(n, H, D), torch.randn(n, HKV, D), torch.randn(n, HKV, D)) for n in lens]
for i, c in enumerate(cached):  # 已缓存的部分预先写进 KV 池
    if c:
        ctx.kv_cache.store_kv(full[i][1][:c], full[i][2][:c], ctx.page_table[i, :c], 0)

batch = Batch(reqs=[req(n, c, i) for i, (n, c) in enumerate(zip(lens, cached))], phase="prefill")
batch.padded_reqs = batch.reqs
batch.out_loc = torch.cat([ctx.page_table[i, c:n] for i, (n, c) in enumerate(zip(lens, cached))])
backend.prepare_metadata(batch)
m = batch.attn_metadata
print("cu_seqlens_q :", m.cu_seqlens_q.tolist(), " (每个请求本轮的 query 起止)")
print("cache_seqlens:", m.cache_seqlens.tolist(), "  (每个请求的 KV 总长)")
print("page_table   :", m.page_table.tolist())
print("out_loc      :", batch.out_loc.tolist())
print("最后位置下标 :", m.get_last_indices(3).tolist())

q = torch.cat([f[0][c:] for f, c in zip(full, cached)])
k = torch.cat([f[1][c:] for f, c in zip(full, cached)])
v = torch.cat([f[2][c:] for f, c in zip(full, cached)])
out = backend.forward(q, k, v, layer_id=0, batch=batch)
ref = torch.cat([
    F.scaled_dot_product_attention(f[0].transpose(0, 1), f[1].transpose(0, 1), f[2].transpose(0, 1),
                                   is_causal=True, scale=D**-0.5, enable_gqa=True).transpose(0, 1)[c:]
    for f, c in zip(full, cached)])
print("与完整因果注意力的最大误差:", (out - ref).abs().max().item())

运行结果

cu_seqlens_q : [0, 6, 9, 10]  (每个请求本轮的 query 起止)
cache_seqlens: [6, 7, 9]   (每个请求的 KV 总长)
page_table   : [[0, 1, 2, 3, 4, 5, 0, 0, 0], [40, 41, 42, 43, 20, 21, 22, 0, 0], [50, 51, 52, 53, 54, 55, 56, 57, 58]]
out_loc      : [0, 1, 2, 3, 4, 5, 20, 21, 22, 58]
最后位置下标 : [5, 8, 9]
与完整因果注意力的最大误差: 5.960464477539062e-07

几点观察：

• 请求 B 的 KV 在池中不连续：前缀缓存命中的 4 个 token 在 40～43，本轮新算的 3 个在 20～22。out_loc 只包含本轮新 token 的位置。
• page table 按 max_seqlen_k = 9 截取，短的请求后面是无关的 0，只要用 cache_seqlens 截断就不会读到。
• get_last_indices 给出 [5, 8, 9]：三个请求最后一个 query 在展平后 q 中的下标。
• 与逐个请求的完整因果注意力相比，误差在 float32 舍入的量级。

创建后端：单个或组合

minisgl/attention/__init__.py · create_attention_backend

def create_attention_backend(backend: str, config: ModelConfig) -> BaseAttnBackend:
    """backend 可以是单个名字，也可以是 "prefill后端,decode后端" 的组合。"""
    validate_attn_backend(backend, allow_auto=False)
    if "," in backend:
        p_name, d_name = backend.split(",", 1)
        if p_name != d_name:
            logger.info(f"Hybrid attention backend: prefill={p_name}, decode={d_name}")
            return HybridBackend(create_attention_backend(p_name, config),
                                 create_attention_backend(d_name, config))
        backend = p_name
    return SUPPORTED_ATTENTION_BACKENDS[backend](config)

后端通过注册表按名字创建。"fa,fi" 表示 prefill 用 FlashAttention、decode 用 FlashInfer，由 HybridBackend 按 batch.is_prefill 转发：

minisgl/attention/base.py · HybridBackend

class HybridBackend(BaseAttnBackend):
    """prefill 和 decode 用不同的后端，例如 Hopper 上 prefill 用 FA3、decode 用 FlashInfer。"""

    def __init__(self, prefill_backend: BaseAttnBackend, decode_backend: BaseAttnBackend) -> None:
        self.prefill_backend = prefill_backend
        self.decode_backend = decode_backend

    def _pick(self, batch: Batch) -> BaseAttnBackend:
        return self.prefill_backend if batch.is_prefill else self.decode_backend

    def forward(self, q, k, v, layer_id: int, batch: Batch):
        return self._pick(batch).forward(q, k, v, layer_id, batch)

    def prepare_metadata(self, batch: Batch) -> None:
        self._pick(batch).prepare_metadata(batch)

    def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None:
        self.decode_backend.init_capture_graph(max_seq_len, bs_list)  # 只有 decode 会被捕获

    def prepare_for_capture(self, batch: Batch) -> None:
        self.decode_backend.prepare_for_capture(batch)

    def prepare_for_replay(self, batch: Batch) -> None:
        self.decode_backend.prepare_for_replay(batch)

CUDA Graph 只捕获 decode，所以三个 graph 相关的方法只转发给 decode 后端。--attn auto 时，引擎在 CPU 上选 torch，在 Hopper（SM90）上选 fa,fi，其他 GPU 上选 fi。注意还有一条 dtype 的约束：FlashInfer 和 FlashAttention 的 kernel 都只为 fp16 / bf16 编译，所以 --dtype float32（本书与 Hugging Face 逐 token 对比时用的精度）在 GPU 上也会退回 torch 后端。

与官方的差异：参考后端

官方没有 PyTorch 后端（SGLang 正式版有一个 torch_native 后端，作用类似）。我们的 torch 后端元数据格式与官方的 FlashAttention 后端相同，CUDA Graph 的三个钩子也按 FlashAttention 后端的方式实现，所以第 18 章可以在 CPU 上验证 CUDA Graph 的数据流。官方在 Blackwell（SM100）上默认选 trtllm 后端，我们没有实现它。

官方实现

• 接口与组合后端：attention/base.py · BaseAttnBackend（第 18～34 行）、attention/base.py · HybridBackend（第 37～63 行）
• 元数据格式参照：attention/fa.py · FlashAttentionBackend.prepare_metadata（第 67～105 行）
• 自动选择后端：engine/engine.py · _adjust_config（第 218～233 行）

测试

tests/test_ch05_attention.py · test_torch_backend_mixed_batch_with_prefix_hits

def test_torch_backend_mixed_batch_with_prefix_hits():
    """两个请求：A 从头 prefill 6 个 token；B 前 4 个 token 已在缓存里，只算后 3 个。"""
    engine = build_engine(QWEN3)
    identity_page_table(engine, 2)
    backend, layer, H, Hkv, D = engine.attn_backend, 5, 16, 8, 128
    torch.manual_seed(0)
    full = {name: (torch.randn(n, H, D), torch.randn(n, Hkv, D), torch.randn(n, Hkv, D))
            for name, n in (("A", 6), ("B", 7))}
    # B 的前 4 个 token 预先写进 KV 池（模拟前缀缓存命中）
    loc_b = engine.page_table[1, :4]
    engine.kv_cache.store_kv(full["B"][1][:4], full["B"][2][:4], loc_b, layer)
    a, b = new_req(list(range(6)), 0), new_req(list(range(7)), 1, cached_len=4)
    from minisgl.core import Batch

    batch = Batch(reqs=[a, b], phase="prefill")
    batch.padded_reqs = batch.reqs
    batch.out_loc = torch.cat([engine.page_table[0, :6], engine.page_table[1, 4:7]])
    backend.prepare_metadata(batch)
    q = torch.cat([full["A"][0], full["B"][0][4:]])
    k = torch.cat([full["A"][1], full["B"][1][4:]])
    v = torch.cat([full["A"][2], full["B"][2][4:]])
    out = backend.forward(q, k, v, layer, batch)
    scale = D**-0.5
    ref_a = dense_causal(*full["A"], scale)
    ref_b = dense_causal(*full["B"], scale)[4:]
    assert torch.allclose(out[:6], ref_a, atol=1e-5) and torch.allclose(out[6:], ref_b, atol=1e-5)
    # LM head 只需要每个请求最后一个位置
    assert batch.attn_metadata.get_last_indices(2).tolist() == [5, 8]
    engine.shutdown()

怎么讲清楚

讲注意力后端：两件事分开——每个 batch 一次 prepare_metadata（描述变长和分页：cu_seqlens_q 划分各请求的 query，cache_seqlens 给出 KV 长度，page table 给出 KV 位置），每层一次 forward（先把本轮的 K、V 写进缓存，再算注意力，否则本轮的 token 看不到自己）。带前缀的 prefill 里因果掩码右下角对齐：KV 长 7、本轮 3 个 query 时，第一个 query 能看到前 5 个 key。元数据每个 batch 只准备一次，因为所有层的 batch 结构都一样。参考后端用 Python 循环 + SDPA 实现，是验证 FlashInfer / FlashAttention 后端的基准；prefill 和 decode 可以用不同的后端。

练习

1. 参考后端对每个请求都构造一个 [n, kv_len] 的掩码。decode 时 n = 1，掩码全为真。给 decode 写一条快路径：把所有请求的 KV 按最大长度补齐成 [bs, max_len, H, D]，用一次 batched SDPA 算完。
2. 如果忘了在 forward 里先调用 store_kv，prefill 的结果会怎样？decode 呢？
3. cu_seqlens_q 在"没有前缀命中的 prefill"时等于什么？在"纯 decode"时等于什么？（提示：看第 17 章 FlashInfer 后端里的两个快路径。）

参考答案

1. 用 cache_seqlens 构造 [bs, 1, max_len] 的长度掩码（arange(max_len) < len），KV 用 page_table[:, :max_len] 一次 gather 出来；注意补齐部分要被掩掉。
2. prefill 时本轮的 key 读到的是池中的旧数据（未初始化或别的请求的），结果错误；decode 时最后一个位置（当前 token）的 KV 是旧的，同样错误。只有"先写后读"，本轮 token 才能看到自己。
3. 没有前缀命中时每个请求的 query 数等于 KV 长度，cu_seqlens_q == cu_seqlens_k；纯 decode 时每个请求 1 个 query，cu_seqlens_q = [0, 1, 2, ..., bs]。

小结

• 注意力后端的两件事：每个 batch 一次 prepare_metadata，每层一次 forward（先写 KV、再算注意力）。
• 元数据描述变长和分页：cu_seqlens_q 划分 query，cache_seqlens 给出 KV 长度，page table 给出 KV 位置；因果掩码右下角对齐。
• 参考后端用 Python 循环 + SDPA 实现，语义清楚、CPU 可跑，是验证 GPU 后端的基准。
• 后端按名字创建，可以 prefill 和 decode 用不同的后端。

本章练习题（1 道）

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

• 189. 注意力后端的元数据（中等）

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