SGLang 源码完全拆解（八）：Attention Backend——统一 metadata 如何落到不同注意力 Kernel

元数据准备：一次编译、多态入口

后端最核心的方法族是 init_forward_metadata 系列。当前源码把它拆成三个层次：

1
class AttentionBackend(ABC):
2
    def init_forward_metadata(self, forward_batch: ForwardBatch):
3
        """Eager entry point. Default = _out_graph(fb) + _in_graph(fb)."""
4
        self.init_forward_metadata_out_graph(forward_batch)
5
        self.init_forward_metadata_in_graph(forward_batch)
6

7
    def init_forward_metadata_out_graph(self, forward_batch, in_capture=False):
8
        """Per-iter metadata prep — runs outside ``with graph.capture():``."""
9

10
    def init_forward_metadata_in_graph(self, forward_batch):
11
        """Graph-recordable static-shape GPU op. Default is no-op."""

这套拆分的动机是 CUDA Graph 与动态 shape 的冲突。元数据准备里既有「必须读主机侧动态信息」的操作，也有「纯静态形状的 GPU 小算子」：前者不能进捕获区域（capture 期间一旦出现 .item()/.cpu()/动态形状的 torch.empty()，重放就会失效或报错），后者放进去反而省掉每次重放前的重复执行。于是契约按执行位置切开：out_graph 方法承载所有主机侧、动态形状、不可捕获的逻辑；in_graph 方法承载纯静态 GPU 算子，在 with graph.capture(): 内被录进图，重放时自动执行。源码在 in_graph 的注释里直接写明了 lint 约束——禁止调用 .item()、.cpu()、.tolist() 与动态形状分配。

类注释里有一条重要记录：legacy 的 init_forward_metadata_capture_cuda_graph 与 init_forward_metadata_replay_cuda_graph 已从 ABC 中彻底移除。旧版后端各写一套捕获/重放钩子，导致捕获分支与重放分支的漂移极其难查；新契约把两者合并成同一个 out_graph(fb, in_capture)，只在 body 确实分叉的后端里通过 in_capture 参数区分。这是一次典型的「用调用时机表达状态，而不是用方法名表达状态」的接口收敛。

前向执行：四种模式的分派

执行入口是 forward(q, k, v, layer, forward_batch, save_kv_cache=True)。它是一个模板方法，按本轮 ForwardMode 分派：

1
def forward(self, q, k, v, layer, forward_batch, save_kv_cache=True, **kwargs):
2
    if forward_batch.forward_mode.is_idle():
3
        return q.new_empty(q.shape[0], layer.tp_q_head_num * layer.v_head_dim)
4
    elif forward_batch.forward_mode.is_decode():
5
        return self.forward_decode(q, k, v, layer, forward_batch,
6
                                   save_kv_cache=save_kv_cache, **kwargs)
7
    elif forward_batch.forward_mode.is_mixed() and is_npu():
8
        return self.forward_mixed(...)
9
    else:
10
        return self.forward_extend(q, k, v, layer, forward_batch,
11
                                   save_kv_cache=save_kv_cache, **kwargs)

四个分支各有一条设计理由：

• idle 分支不执行任何 kernel。上一章讲过 idle 批次是重叠调度与投机解码产生的空转轮次：GPU 上这批「没有真实 token」的请求仍然要参与 collective，但注意力层没有有效输入。此时唯一正确的行为是返回一个形状正确的空输出继续往下走，而不是让 kernel 收到零长度张量。返回缓冲的形状由 tp_q_head_num 与 v_head_dim 乘出来，与真实路径严格一致。
• decode 与 extend 是两个不同方法，不是同一个方法加参数。因为两者需要的 kernel 结构根本不同：decode 每请求一个查询 token、要读整段历史 KV，主流方案是分页 decode kernel 加 KV 切分（split-KV）；extend 是变长多 token 查询，主流方案是变长 prefill kernel 加因果 mask 或显式 mask。让一个方法内 if decode 分支，只会让每个后端的主逻辑被揉成一团。
• is_mixed() 是 NPU 专用分支。昇腾平台的硬件路径需要混合模式单独处理，CUDA 平台走不到这里。

forward_decode、forward_extend、forward_mixed 在基类里都是 NotImplementedError，由具体后端实现。基类只固定「什么时候调用谁」，不假设「怎么算」。

CUDA Graph 状态与填充值

后端必须实现 init_cuda_graph_state(max_bs, max_num_tokens)：按最大批次预分配那些内容每轮变化、但地址必须固定的缓冲（例如 FlashInfer 的 cuda_graph_kv_indices、Triton 的按 bs 切分的 kv_indptr 视图）。与之配套的 get_cuda_graph_seq_len_fill_value() 决定捕获时序列长度张量填什么值——FlashInfer 与 Triton 的实现都返回 1。这个 1 不是随手选的：捕获阶段 kernel 会真实读一遍 seq_lens，填 0 会让某些 kernel 进入零长度退化路径、甚至触发除零；填一个小的合法值可以让捕获走完整流程，真正的内容在重放前由 out_graph 阶段刷新。

辅助能力：不强迫所有后端实现全部功能

基类还声明了一批能力标志，它们的价值在于让上层可以问而不必猜：

还有一个容易被忽略的小枚举 SharedReadEnds。它声明「本后端何时读完调度器共享的数据」：

1
class SharedReadEnds(Enum):
2
    PRE_REPLAY = 1   # init_forward_metadata_out_graph 之后读完
3
    IN_REPLAY = 2    # init_forward_metadata_in_graph 之后读完
4
    POST_REPLAY = 3  # 未实现元数据快照
5
    UNKNOWN = 4      # 未审计 -> 保守的整前向栅栏

默认实现给 decode 与 target_verify 返回 IN_REPLAY，其余模式返回 UNKNOWN（保守取整前向栅栏）。这个声明服务于重叠调度：CPU 侧在准备下一轮调度时需要知道「上一轮的后端元数据构建是否已经读完共享张量」，晚读的部分要求更宽的屏障。把「读过哪些共享数据」显式声明出来，比让每个调用点猜一张全局保守锁便宜得多。这类微小的枚举正是「后端不仅执行 kernel、还参与时序正确性证明」的证据。

注册：一个装饰器与一张字典

attention_registry.py 的注册表简单到极致：

1
ATTENTION_BACKENDS = {}
2

3
def register_attention_backend(name):
4
    def decorator(fn):
5
        ATTENTION_BACKENDS[name] = fn
6
        return fn
7
    return decorator
8

9
@register_attention_backend("flashinfer")
10
def create_flashinfer_backend(runner):
11
    if not runner.use_mla_backend:
12
        from ...flashinfer_backend import FlashInferAttnBackend
13
        return FlashInferAttnBackend(runner, init_new_workspace=runner.init_new_workspace)
14
    else:
15
        from ...flashinfer_mla_backend import FlashInferMLAAttnBackend
16
        return FlashInferMLAAttnBackend(runner)

两个细节值得展开。第一，工厂函数内部做延迟导入：flashinfer 这个条目在导入注册表模块时就被登记，但只有真正被选中时才 import 对应模块。这让一台不装 FlashInfer 的机器照样能启动别的后端，也让 SSM（Mamba）这类特殊模型不去加载与它无关的注意力依赖。第二，工厂内部按 use_mla_backend 二次分派——MLA 模型的 KV 是压缩潜向量而非标准 K/V 头，同一名字 flashinfer 在 MLA 模型上落到完全不同的实现类。

当前注册表覆盖的后端可以按用途分组：

• 通用 CUDA 后端：flashinfer（分页/ragged 双 wrapper）、fa3/fa4（FlashAttention-3/4）、triton（自研 Triton kernel）、trtllm_mha（TensorRT-LLM 的 MHA kernel）、hpc_ops、torch_native（参考实现）、flex_attention（torch 的 FlexAttention）；
• MLA 专用后端：flashmla、trtllm_mla、tokenspeed_mla、cutedsl_mla、flashinfer_mla；注册表里多条 MLA 条目各有取向，例如 trtllm_mla 的工厂会在「解码上下文并行加投机解码」的组合上直接拒绝启动并提示换 cutedsl_mla 或 tokenspeed_mla——它不转发跨 rank 合并所需的 rank-local LSE；
• 稀疏注意力后端：dsa（DeepSeek Sparse Attention，旧的 nsa 名字以 DeprecationWarning 转发到 dsa）、qsa（Qwen 系稀疏注意力）、dsv4（DeepSeek-V4 压缩注意力，NPU/HIP/CUDA 三个平台分支）；
• 平台后端：aiter（ROCm）、ascend（昇腾）、intel_amx（Intel CPU）、wave。

设计取舍很清楚：后端颗粒度是「一套 kernel 家族」，不是「一个 kernel」。同一家族里 prefill 与 decode 用不同 kernel 是常态，家族内部怎么切由自己决定；注册表只保证「这个名字能被解析到、能被构建」这一件事。

解析：prefill 与 decode 可以是两个不同后端

装配链第二步在 attention_backend_setup.py。resolve_attention_backend_strs 返回一对字符串而不是一个：

1
class ResolvedAttentionBackendStr(msgspec.Struct, frozen=True, kw_only=True):
2
    prefill: str
3
    decode: str
4
    is_draft_override: bool = False
5

6
def resolve_attention_backend_strs(*, model_runner) -> ResolvedAttentionBackendStr:
7
    if model_runner.is_draft_worker and model_runner.draft_attention_backend:
8
        # 草稿 worker 有独立后端，目标与草稿共存于同一进程
9
        return ResolvedAttentionBackendStr(prefill=..., decode=..., is_draft_override=True)
10
    prefill, decode = attention_backends()
11
    return ResolvedAttentionBackendStr(prefill=prefill, decode=decode)

「prefill 与 decode 用不同后端」是这一层最有信息量的设计。prefill 是计算密集的变长多 token 计算，追求矩阵乘利用率；decode 是带宽密集的逐 token 计算，追求 KV 读取效率与切分调度。同一个模型家族里，最优 kernel 往往不同——例如 MLA decode 上 flashmla/trtllm_mla 的块表路径显著优于通用分页 kernel，而 prefill 上 FlashInfer 的 ragged wrapper 通常更合适。允许两者分开配置，服务就能取各自最优解，而不必被单一后端的最短板限制。

草稿 worker 的覆盖逻辑则揭示了另一个约束：投机解码的草稿模型与目标模型共用一个进程，所以草稿的后端不能从进程级配置读——它必须挂在 ModelRunner.draft_attention_backend 上，由草稿侧自己的参数决定。

构建：单后端、双后端、包装器

build_attention_backends 是装配的终点，分支序列本身就是一份设计文档：

1
def build_attention_backends(*, model_runner) -> AttentionBackends:
2
    if model_runner.model_config.attention_arch == AttentionArch.SSM:
3
        mamba_backend = Mamba2AttnBackend(model_runner)   # 纯 SSM：没有注意力
4
        return AttentionBackends(attn_backend=mamba_backend, ...)
5

6
    if get_disagg().enable_pdmux:
7
        # PD 复用：一个 prefill 后端 + sm_group_num 组 decode 后端实例
8
        attn_backend = _build_resolved_backend(..., init_new_workspace=True)
9
        decode_attn_backend_group = [_build_resolved_backend(..., init_new_workspace=False)
10
                                     for _ in range(get_disagg().sm_group_num)]
11
    elif get_exec().overlap.enable_two_batch_overlap and not model_runner.is_draft_worker:
12
        attn_backend = TboAttnBackend.init_new(
13
            lambda: _build_resolved_backend(..., init_new_workspace=False))
14
    else:
15
        attn_backend = _build_resolved_backend(..., init_new_workspace=False)

四个分支对应四种部署形态：

• 纯 SSM 模型（如 Mamba2 架构）根本没有 softmax 注意力层，直接返回 Mamba2AttnBackend，prefill/decode 字符串都记作 mamba2；
• PDmux（prefill/decode 复用同一组 GPU 的多流模式）为每个流组构建一个独立的 decode 后端实例，只有第一个分到新 workspace——多实例共享显存但必须各自持有元数据；
• 双批次重叠（TBO）把批次切给两个子后端（TboAttnBackend 持有 children=[child0, child1]），两份子批次在各自动的 metadata 上并行推进，再回写合并；
• 默认路径：单个后端实例。

构建入口 _build_resolved_backend 处理 prefill ≠ decode 的情形——它把这对比作「两个完整注意力后端」再套一层 HybridAttnBackend：

1
class HybridAttnBackend(AttentionBackend):
2
    """Support different backends for prefill and decode."""
3

4
    def _select_backend(self, forward_mode: ForwardMode) -> AttentionBackend:
5
        if forward_mode.is_decode_or_idle():
6
            return self.decode_backend
7
        elif forward_mode.is_target_verify():
8
            return (self.decode_backend if self.spec_attn_is_decode
9
                    else self.prefill_backend)
10
        else:
11
            return self.prefill_backend

路由规则里唯一的非平凡点在 target_verify：投机解码的验证阶段既像 prefill（多 token 查询）又像 decode（要按 KV 长程读取），所以它跟着 --speculative-attention-mode 的取值路由——该参数声明「验证在这个框架里按哪类负载对待」。这不是猜测，而是配置与后端能力的显式对齐。

HybridAttnBackend 还要处理一批「问哪一半」的转发：init_cuda_graph_state 默认只初始化 decode 侧，只有 DLLM 或「验证走 prefill」的组合才追加初始化 prefill 侧；verify_mask 从验证执行者转发；needs_cpu_seq_lens 按实际会服务的模式 OR。这些转发看起来琐碎，但每一条都避免了一类真实 bug——比如若把不服务验证的 prefill 后端的 needs_cpu_seq_lens 也算进去，每轮 decode 都会被强制多一次 D2H 同步。

包装器：混合模型与线性注意力

attn_backend_wrapper 是最后一道包装，处理的是「同一个模型里既有 softmax 注意力层、又有别的层」的混合架构。当前源码里的分支包括：

• 混合 GDN（Gated DeltaNet，如 Qwen 系混合模型）与混合 KDA/Lightning/Mamba2：完整注意力层用选定的全注意力后端，线性注意力层用对应的线性后端，二者由 HybridLinearAttnBackend 组合；shared 的混合层路由与 mamba 状态管理都在这一层；
• MiniMax 稀疏注意力：完整层与稀疏层各持一个后端，由 MiniMaxHybridAttnBackend 按 sparse_layer_ids 分发；
• Inkling 短卷积混合、Dots3 等特定模型族：各自把「侧车（sidecar）状态元数据」挂进包装对象。

包装器存在的意义是让基座注意力后端对混合架构保持无知。FlashInfer、Triton 的实现不需要知道「这个模型第 12 层是线性注意力」；包装对象在层与后端之间按层号分流，保持内核侧代码对所有模型族通用。

装配完成后，ModelRunner.init_attention_backends 还会做两件收尾：把解析出的 per-mode 字符串盖章到后端对象（prefill_attention_backend_str / decode_attention_backend_str，供模型侧按模式查询），以及让 kv_index_translator.bind_and_verify_backends([attn_backend, decode_attn_backend]) 校验每个后端都能通过翻译器读统一内存池。注释里还记录了一条顺序约束：configure_aux_hidden_state_capture（EAGLE/DFLASH 的辅助隐状态捕获配置）必须在本函数里、CUDA Graph 捕获之前完成——否则捕获出的图不含辅助隐状态输出路径，投机解码会在图里拿到空张量。

一个具体的问题：kernel 只认识张量，批次是 Python 对象

注意力 kernel 的输入是纯张量：查询张量、KV 缓存张量、一串描述「每行查询该读哪些 KV 页」的索引。而 ForwardBatch 是另一套语言：req_pool_indices 是批内第 $i$ 个请求在请求池里的槽号，req_to_token 是「请求槽 × 逻辑位置 → KV 位置」的二维映射，seq_lens 是本轮每请求的有效长度。后端的元数据构造就是这两套语言之间的编译器。

编译出来的第一层结构是分页读地址。FlashInfer 的 decode 元数据由三件套构成，源码里它们的初始化是一个通用小函数：

1
def _init_kv_indices(...):  # 结构示意，合并自 FlashInferIndicesUpdaterDecode
2
    bs = len(paged_kernel_lens)
3
    kv_indptr[1 : bs + 1] = torch.cumsum(paged_kernel_lens, dim=0)
4
    kv_indptr = kv_indptr[: bs + 1]
5
    if wrapper.is_cuda_graph_enabled:
6
        kv_indices = wrapper._paged_kv_indices_buf   # 直接写进图输入缓冲
7
    else:
8
        kv_indices = torch.empty(paged_kernel_lens_sum, dtype=torch.int32, device="cuda")
9
    translator.pack_read_stream(plan, req_pool_indices=..., seq_lens=paged_kernel_lens,
10
                                indptr=kv_indptr, out=kv_indices,
11
                                kv_start_idx=kv_start_idx, kind=read_kind)

三个张量各司其职：

• kv_indptr：长度 $bs+1$ 的前缀和数组，第 $i$ 项是前 $i$ 个请求的 KV 页总数。它是 CSR（压缩稀疏行）格式的行指针——第 $i$ 个请求的页列表切片在 kv_indices[kv_indptr[i]:kv_indptr[i+1]]。用前缀和而不是「每请求一个等长列表」，是因为请求长度不同，二维张量会浪费（padding）或非法（锯齿）；
• kv_indices：长度等于本轮 KV 页总数的扁平列表，内容是该请求读到的物理页号。注意它不是「逻辑页号」，而是经过翻译器落到真实内存池的物理页号；kv_index_translator 的存在正是为了让「逻辑页 → 物理页」的翻译可以按内存池形态（统一池、SWA 池、卸载池）变化，而 kernel 侧不必知道；
• kv_last_page_len：每个请求最后一页有几个有效 token。分页存储的最后一页通常填不满，decode kernel 必须知道「读到最后一页的哪个位置为止」。

最后 wrapper.begin_forward(kv_indptr, kv_indices, kv_last_page_len, num_qo_heads, num_kv_heads, head_dim, page_size=1, ...) 把三件套交给 FlashInfer 的 plan()：这一步在主机侧完成工作划分与负载均衡计算，把 kernel 启动参数固化到 wrapper 内部。真正的 wrapper.forward(q, kv_cache, sm_scale, logits_soft_cap, k_scale, v_scale) 才是逐层执行。

分页 KV 的读地址结构可以用 FlashInfer 论文的块稀疏表（BSR）图来对照理解：当页表按 $(B_r, B_c) = (4, 1)$ 分组时，「每行 4 个查询读 1 个 KV 块的块级视图」正好是 kv_indptr/kv_indices 逻辑的块粒度版本。

图源：FlashInfer 论文（arXiv:2501.01005） Figure 2。SGLang 后端构造的是逐 token 粒度的 CSR 页表（等价于 $B_r = B_c = 1$ 的极端情形）；当某层使用块稀疏注意力时，同样的翻译思路会提升到块粒度。

为什么逐层执行、元数据却每轮只建一次

回到 ModelRunner 的一轮循环：init_forward_metadata(fb) 在模型前向之前被调用一次；随后模型里的 $L$ 个注意力层各自调用 backend.forward(...)，每次从 self.forward_metadata 里取出当轮对象。这个「一次编译、$L$ 次执行」的划分是必须的：元数据构造要读 seq_lens_cpu、做前缀和、可能触发主机侧同步，成本随层数乘起来会很可观；而逐层执行里唯一必须逐层变化的量是 layer_id（选 KV 缓冲）与层的形状参数。

也正因为元数据是批次级对象，FlashInfer 后端才需要 _get_wrapper_idx(layer) 这层间接：一个批次里可能有多类注意力层——全注意力和滑窗层需要不同的 wrapper（页表里 KV 长度上限不同），交叉注意力层又需要别的配置。源码用两种 dispatch 原因区分：

1
class WrapperDispatch(Enum):
2
    SLIDING_WINDOW = ...
3
    CROSS_ATTENTION = ...
4

5
def _get_wrapper_idx(self, layer: RadixAttention):
6
    if self.num_wrappers == 1:
7
        return 0
8
    if self.dispatch_reason == WrapperDispatch.SLIDING_WINDOW:
9
        return layer.sliding_window_size == -1
10
    if self.dispatch_reason == WrapperDispatch.CROSS_ATTENTION:
11
        return layer.is_cross_attention

即：批次级元数据对象里放一组 wrapper，逐层执行时按层的属性选索引。滑窗层的 KV 读取被限制在窗口内，cross-attention 层（编码器-解码器模型）读的是编码器侧 KV，这两类层若共用同一份页表就会读到错误数据。把「哪些层需要分开的元数据」做成批次级枚举、把「这一层属于哪一类」做成层级布尔判断，是这套设计里最小的实现方案。

这里有一个容易被忽略的工程细节，直接写在后端构造的注释里：FlashInfer 把 num_qo_heads 烘焙进 plan，任何「跑得比 plan 时假设的头数更多」的层都会被静默算错。因此 plan 时采用「按批内最大头数过度规划」的策略——过度规划是安全的，欠规划不是。GQA/MQA 模型里不同层查询头数可能不同，这个细节决定了后端必须按最坏情况准备。

图源：FlashInfer 论文（arXiv:2501.01005） Figure 11。GQA 里多个查询头共享一组 KV 头，kernel 侧可以把头组融合进同一维做并行；而框架侧必须保证 plan 的头数参数不小于任何一层的实际头数。

extend 的两种形态：ragged 直算与 paged + merge

prefill/extend 的元数据构造比 decode 多一层选择。源码先决定本轮走哪条路径：

1
use_ragged = (not self.enable_deterministic
2
              and not is_in_tc_piecewise_cuda_graph()
3
              and not self.use_paged)
4
extend_no_prefix = not any(forward_batch.extend_prefix_lens_cpu)

• ragged 直算：当批次里所有请求都没有命中缓存前缀（extend_no_prefix 为真）时，整个批次是一次干净的变长 prefill——所有 token 都是「新的」，直接调用 ragged wrapper：forward(q, k, v, causal=..., sm_scale=..., logits_soft_cap=...)，KV 也在这次调用里由 save_kv_cache 写入页表。
• paged + merge：当有请求命中了前缀缓存，本轮 extend 的每个请求分成两段——已缓存在页表里的前缀段（$p_i$ 个 token）与需要新计算的后缀段（$e_i$ 个 token）。前缀段的 KV 在内存池里（分页），后缀段的 KV 在本次计算的 k/v 张量里（连续）。FlashInfer 的做法是分别算两段、再合并注意力状态：

其中 $o$ 是未归一化的注意力输出、$s$ 是对应的 log-sum-exp 统计量。合并两个 online-softmax 状态的标准形式是：

源码里对应的调用是：

1
o1, s1 = self.prefill_wrapper_ragged.forward_return_lse(q, k, v, causal=..., ...)
2
o2, s2 = prefill_wrapper_paged.forward_return_lse(q, kv_cache, causal=False, ...)
3
o, _ = _safe_merge_state(o1, s1, o2, s2)

三个设计点：其一，ragged 段传 causal=True（段内因果），paged 段传 causal=False（前缀段全部可见，因果性已经由「前缀在时间上更早」这一事实保证）；其二，forward_return_lse 返回统计量而不是只返回结果，是为了让合并是数值精确的在线 softmax 合并，而不是近似加权；其三，_safe_merge_state 的「safe」有实义——源码在前面计算了「合并阶段最多能安全处理多少个头」的上界（按 head_dim 与元素大小推算的 _merge_state_max_safe_num_heads），超限时走分块合并，避免 merge kernel 的寄存器压力爆掉。

mask 从哪里来：因果位、滑窗、验证掩码

「构造 KV 索引与 mask」里的 mask 部分同样属于后端。与一些实现「每轮物化一张完整注意力掩码张量」的做法不同，SGLang 的 mask 有三种形态，按需要才升级：

第一形态是因果位。 绝大多数 extend/decode 的可见性完全由位置决定，所以后端根本不构造掩码张量，只把 causal 布尔值传给 kernel——kernel 内部按查询位置与 KV 长度推导逐格可见性。前面 extend 合并路径里 ragged 段传 causal=True、paged 前缀段传 causal=False 就是这一形态的用法：前缀段在时间上必然早于所有查询，因果性自然成立，无需逐格掩码。

第二形态是滑窗。 层带 sliding_window_size 时，后端的处理是把 window_left=layer.sliding_window_size 作为参数交给 kernel，由 kernel 裁剪可见范围。但滑窗与「paged 前缀 + ragged 后缀」合流时，单纯参数不够了：前缀段在页表里的绝对位置与查询的绝对位置之间需要逐格判断「是否落在窗口内」。FlashInfer 后端为此实现了 _build_swa_prefix_custom_mask——对每个请求，令查询的绝对位置为 prefix_len .. prefix_len + extend_len、KV 的绝对位置为 kv_start .. prefix_len，逐格保留 $k \ge q - w$ 的位置；只有当存在越窗格子时才真正构造掩码（全在窗内时返回 None 走快路径）。注意这段代码先把三组长度张量转成主机列表做判断——这是后端里少见的显式主机侧分析，也说明掩码构造的成本只有在必要时才值得付。

第三形态是投机解码的验证掩码。 树式草稿（EAGLE 系）的验证阶段，可见性与树结构有关，无法用因果位或滑窗表达，必须显式给出。源码里两个对象配合完成：

1
class VerifyMask(msgspec.Struct, frozen=True):
2
    """The target-verify mask."""
3
    buffer: torch.Tensor        # 草稿阶段原地写入的掩码缓冲
4
    mode: TreeMaskMode
5
    max_bs: int
6
    is_read: bool = True
7

8
    def fits(self, bs: int) -> bool: ...
9

10
def tree_mask_numel(mode, bs, num_draft_tokens, max_context_len) -> int:
11
    if mode == TreeMaskMode.QLEN_ONLY:
12
        per_req = num_draft_tokens * num_draft_tokens          # KB 级
13
    elif mode == TreeMaskMode.FULL_MASK:
14
        per_req = num_draft_tokens * (max_context_len + num_draft_tokens)  # 长上下文下上百 MB

设计要点：掩码由草稿阶段通过 build_tree_kernel_efficient 原地写入缓冲，目标验证轮直接读取——这让验证 worker 省掉一次 seq_lens_sum 的 D2H 同步；缓冲尺寸按 max_bs 一次性分配，fits() 在批次超界时提前暴露；mode 与缓冲必须成套传递（注释原话：只拿缓冲不拿布局，kernel 会写出读方不认识的形状）；结构体是 frozen 的，扩容只能整体替换，防止 max_bs 与缓冲脱节。掩码长度则通过 mask_indptr 表达——每请求的掩码格数是 num_tokens_per_req × (seq_lens + num_tokens_per_req)，前缀和之后交给 kernel 做 CSR 切片，与 KV 页表的 kv_indptr 完全同构。

为什么不做成「统一物化」？回顾三种形态的成本：因果位零成本；滑窗在必要时才构造、且按需裁剪；验证掩码必须显式但由草稿阶段顺路写好。如果统一成「每轮构造完整掩码」，decode 这种 $bs \times 1 \times L$ 的场景会白白付出 $bs \times L$ 的内存与构造时间，长上下文下完全不可接受。掩码的形态选择本身就是一种优化，这也是后端契约把 mask 参数（causal、window_left、custom_mask/mask_indptr）全部暴露在逐层调用签名上的原因。

其他后端的元数据：同一个契约，不同物化

Triton 后端的 ForwardMetadata 是一个 dataclass，字段列表本身就是它的执行计划：

1
class ForwardMetadata:
2
    attn_logits: torch.Tensor      # (bs, num_head, max_kv_splits, v_head_dim) 部分结果
3
    attn_lse: torch.Tensor         # (bs, num_head, max_kv_splits) 部分统计量
4
    max_extend_len: int
5
    num_kv_splits: torch.Tensor    # 每请求的实际切分数
6
    kv_indptr / kv_indices         # decode/verify 用分页页表
7
    qo_indptr                      # extend 的查询 CSR
8
    custom_mask / mask_indptr      # 验证阶段的显式 mask
9
    window_kv_indptr / window_kv_indices / window_num_kv_splits  # 滑窗层专用
10
    ...

它与 FlashInfer 的差异在把 decode 的 KV 切分暴露到了元数据层。Triton decode kernel 采用 split-KV：把每个请求的长 KV 切成 num_kv_splits[i] 段并行归约，每段产出部分结果 attn_logits 与部分统计 attn_lse，最后由 reduction kernel 合并。切分数由 get_num_kv_splits 按序列长度动态选择：序列越长切得越多，用更多并行度对抗单序列的长尾。FlashInfer 的等价机制藏在它自己的 plan() 内部（fixed_split_size/disable_split_kv 参数），而 Triton 后端把它显式为框架可见的元数据字段——这解释了为什么 Triton 后端需要为「lean decode」维护额外的持久化网格部分结果缓冲（lean_Mp/lean_Lp/lean_Op/lean_locks）。

FlashInfer 论文的负载均衡调度图说明了这类「切分 + 均匀分发」的一般问题：变长序列在 CTA 粒度上的工作负载极不均匀，调度器需要按序列长度信息把工作量摊平，且必须与 CUDA Graph 的静态配置兼容。

图源：FlashInfer 论文（arXiv:2501.01005） Figure 6。SGLang 的 Triton 后端用 num_kv_splits 做序列内切分，FlashMLA 用 get_mla_metadata 产出的 tile 调度元数据做跨序列分发，FlashInfer 在自己的 plan 阶段完成同类摊平——三者是同一问题的三种解法。

FlashMLA 后端（MLA 模型的 decode 主力，继承自 FlashInferMLAAttnBackend）的元数据是三件不同的东西：

1
mla_metadata, num_splits = get_mla_metadata(
2
    forward_batch.seq_lens.to(torch.int32), self.num_q_heads, 1,
3
    is_fp8_kvcache=self.is_fp8_kvcache,
4
)
5
self.forward_metadata = FlashMLADecodeMetadata(mla_metadata, num_splits, block_kv_indices)

• 块表 block_kv_indices：形状是（批次, 每序列的块数），块大小固定 64 token。构造路径有两条——统一内存池下用 kv_index_translator.copy_page_table(plan, out=block_kv_indices) 从逻辑视图直接拷贝；否则用 create_flashmla_kv_indices_triton[(bs, 块数)] 从 req_to_token 与 req_pool_indices 现场生成。第二条路径正是上一节「请求槽映射表 → 物理页表」翻译的 Triton 版本；
• tile_scheduler_metadata 与 num_splits：由 get_mla_metadata 产出，是 kernel 内部的 tile 级调度计划——把「哪些块的 KV 归哪个 CTA、跨序列怎么摊平、每个序列切几份」预先算好，kernel 只按表执行。它和 FlashInfer 的 plan() 是同一思想：把调度计算移出 kernel 主循环，换取主循环里的纯数据搬运；
• seq_lens_k（仅 draft-extend 模式）：KV 侧的有效长度，与查询窗口长度分开。

调用形态很直白：flash_mla_with_kvcache(q, k_cache, block_table, cache_seqlens, tile_scheduler_metadata=..., num_splits=...)。MLA 的 KV 是压缩潜向量（每个 token 只有几百字节），块表是它最主要的寻址结构；这也是为什么 MLA 的 decode 路径不套用通用分页三件套——潜向量的宽行形态让「页级 indptr」反而不如固定 64-token 块表划算。

一个有趣的细节：FlashMLA 后端的 init_forward_metadata 只处理 decode/verify/draft-extend 三类模式，其余模式回落到父类（super().init_forward_metadata(forward_batch)，即 FlashInfer MLA 的通用实现）。也就是说，同一后端对象内部也允许「decode 用专用 kernel、extend 用通用 kernel」的组合——与 HybridAttnBackend 的跨后端组合是同一种思想，只是粒度退到了类内部。

第一件：预分配地址固定、内容可变的缓冲

CUDA Graph 捕获后，图内 kernel 引用的张量地址被冻结。分页索引每轮都变，所以正确的做法是预分配足够大的缓冲，每轮只刷新内容。FlashInfer 的 init_cuda_graph_state 就干这件事：

1
def init_cuda_graph_state(self, max_bs, max_num_tokens, kv_indices_buf=None):
2
    self.cuda_graph_kv_indices = [torch.zeros(最大页数, dtype=torch.int32, device=...)]
3
    self.cuda_graph_kv_indices += [x.clone() for _ in range(self.num_wrappers - 1)]
4
    # wrapper 组里每个 wrapper 一份，避免共享导致互相踩踏
5

6
def _create_decode_wrappers(self, bs, num_tokens):
7
    return [DecodeWrappers(int_workspace_buffer=..., page_size=1,
8
                           paged_kv_indptr_buffer=self.kv_indptr[i][: num_tokens + 1],
9
                           paged_kv_indices_buffer=self.cuda_graph_kv_indices[i],
10
                           paged_kv_last_page_len_buffer=self.kv_last_page_len[:num_tokens], ...)]

要点有三：缓冲按 max_bs/max_num_tokens 上界分配，重放时按当前真实 bs 取切片视图（[:num_tokens+1]），切片与原缓冲共享地址；每个 wrapper 一份独立缓冲，因为滑窗层与全注意力层若共用一个 kv_indices 缓冲，捕获出的两个 kernel 会互相覆盖；缓冲初始化为 0（第 0 项按约定归零），因为捕获阶段会真实执行一遍。

捕获时哪些内容应该清零还有一条具体规则：捕获阶段把 cuda_graph_swa_out_cache_loc[:n] 清零、重放阶段才 copy 真实值。原因是捕获时还没有真实批次，缓冲里残留的旧页号会被录进图并被 kernel 读到（用于写 KV 的定位），必须清成一个安全值。

第二件：分离「捕获时算一次」与「每次重放前刷新」

init_forward_metadata_out_graph(fb, in_capture) 的 in_capture 参数就为此存在。捕获时（in_capture=True）需要额外做几类事：

• 调用 _prepare_cuda_graph_metadata(bs, num_tokens, forward_mode, spec_info) 为这个 bs 建好 wrapper 组并存入 decode_cuda_graph_metadata[bs]；
• 把 begin_forward 换成无同步版本。源码在捕获时执行：

1
if in_capture and forward_mode.is_decode_or_idle():
2
    for w in self.decode_cuda_graph_metadata[bs]:
3
        w.begin_forward = partial(fast_decode_plan, w)

fast_decode_plan 的意义：标准的 FlashInfer plan() 需要读主机侧的序列长度信息，会触发多次 D2H 同步；而在重放路径上，这些信息其实已经在捕获时确定（或由重放前的 host 侧调用提供），每次重放再同步一遍会把 CPU 卡在正在执行的图后面。注释里对 DFLASH target-verify 那条分支写得更直白：换成 sync-free 的 plan 可以「删除每次重放的三个阻塞式 .to("cpu") 读取」。这是 CUDA Graph 与主机侧调度耦合处最典型的性能陷阱——图本身没有同步，但准备元数据的代码有。

• 把 SWA 写目标缓冲绑定到元数据上（self.forward_metadata.swa_out_cache_loc = self.cuda_graph_swa_out_cache_loc[:n]），保证图内 kernel 引用的地址稳定。

第三件：声明元数据是否被图跨断点引用

上一章的 Breakable CUDA Graph（BCG）需要把 eager 段插进捕获过程，此时一个后端可能把自己的元数据张量交给 kernel 跨图断点使用。契约为此留了两个钩子与一组标志：

1
use_captured_forward_metadata_for_breakable_cuda_graph: bool = False
2
prefill_cuda_graph_max_batch_size: Optional[int] = None
3

4
def init_forward_metadata_for_breakable_cuda_graph_capture(self, forward_batch):
5
    """Create forward metadata whose tensor addresses will be graph-captured."""
6
    raise NotImplementedError()
7

8
def prepare_forward_metadata_for_breakable_cuda_graph_replay(
9
    self, capture_metadata, forward_batch, *, static_forward_batch=None) -> None:
10
    """Refresh captured metadata for the current batch before BCG replay."""
11
    raise NotImplementedError()

语义是：选择加入（opt-in）的后端在捕获阶段创建元数据对象（地址被录进图），在每次重放前原位刷新其动态字段。选择不加入的后端则每次重放前重建元数据（地址本来就不进图）。前者省掉每轮的对象构造，代价是后端必须保证「所有被 kernel 读的字段都在刷新范围内」；后者更简单但更贵。把这个选择交给后端自己声明，是因为只有后端知道自己的元数据张量有哪些会被 kernel 跨断点读到。

与之配套的还有 SharedReadEnds 声明（前面讲过的 PRE_REPLAY/IN_REPLAY/UNKNOWN）：BCG 在发布「重放前事件」之前，需要知道后端是否已经读完调度器共享的数据，才能安全地让 CPU 侧继续准备下一轮。没有这个声明，BCG 就只能对每个后端都套整前向的粗粒度栅栏，流水线收益归零。