FlashInfer GDN Prefill 完全指南：chunk_gated_delta_rule API 与 Cake GDN 上下文并行后端

FlashInfer 的 Gated Delta-Rule（GDN）线性注意力系列包含完整的 prefill 侧内核，chunk_gated_delta_rule 是其中负责「chunked GDN scan」的入口：它在 decode 循环接管之前，用一整段 prompt 初始化 recurrent state。本文以 docs/api/gdn_prefill.rst 为骨架，结合 flashinfer/gdn_prefill.py 的调度器实现、flashinfer/jit/cake_gdn_cp_backend.py 的 JIT 注册表以及 tests/gdn/test_cake_gdn_cp_backend.py 的测试用例，系统讲解该 API 的全部参数语义、状态池与检查点机制，以及 SM100a/SM103a 上 Cake GDN 上下文并行（Context Parallel，CP）后端的设计与用法。

读完本文，你将掌握：如何用 chunk_gated_delta_rule 完成 GDN prefill、如何通过 state_indices 接入页式状态池、如何配置中间状态检查点，以及如何在 Blackwell 架构上显式启用 Cake GDN CP 后端并获得 CUDA Graph 友好的调度行为。

GDN prefill：decode 之前的 recurrent state 初始化

在 Gated Delta Rule 线性注意力中，序列状态是一个按 token 逐步更新的递推对象。prefill 阶段的任务是把整段 prompt 一次性「扫描」进状态：

• chunk_gated_delta_rule 是一个 chunked GDN scan——它把每个序列按 64 token 的 chunk 切分，逐 chunk 推进递推，从而初始化 recurrent state，供后续 decode 阶段接手。
• 模块层面，flashinfer.gdn_prefill 是 GDN 的 prefill 侧内核集合，模块只导出一个公共入口 chunk_gated_delta_rule，而 decode 侧的对应内核位于 flashinfer/gdn_kernels（如 gated_delta_rule、gated_delta_rule_mtp 等）。
• 该内核同时支持 GQA（num_q_heads > num_k_heads = num_v_heads）与 GVA（num_v_heads > num_q_heads = num_k_heads）两种 head 映射，最终状态布局为 [N, H, V, K]。

从 flashinfer/gdn_prefill.py 的调度器看，chunk_gated_delta_rule 是唯一公共入口，内部按计算能力（SM90/SM100/SM103/SM120）分派到不同 CuTe-DSL 内核或 Cake 后端。

chunk_gated_delta_rule API 全参数解析

函数签名如下（来自 flashinfer/gdn_prefill.py）：

chunk_gated_delta_rule(
    q, k, v,
    g=None, beta=None, scale=None,
    initial_state=None, output_final_state=False,
    cu_seqlens=None, use_qk_l2norm_in_kernel=False,
    output=None, output_state=None,
    state_checkpoints=None, checkpoint_cu_starts=None,
    checkpoint_every_n_tokens=0,
    use_cp="auto", state_indices=None,
    _cp_chunk_len=None,
    backend="auto", max_seqlen=None,
)

各参数语义汇总如下：

参数	类型/形状	说明
q / k / v	[total_seq_len, num_heads, head_size]	查询/键/值，必须连续且位于 CUDA 上；SM100 路径要求 head_size == 128
g	可选 [total_seq_len, num_sab_heads] float32	遗忘门 alpha；None 时默认为全 1；num_sab_heads = max(num_q_heads, num_v_heads)
beta	可选 [total_seq_len, num_sab_heads] float32	更新门 beta；None 时默认为全 1
scale	可选 float	注意力分数缩放；None 时默认为 1 / sqrt(head_size)
initial_state	可选 [num_seqs, num_sab_heads, head_size, head_size]	初始 KV 状态；dtype 支持 float32/bfloat16/float16/float8_e4m3fn/float8_e5m2；None 从零状态开始。给出 state_indices 时是状态池 [N_pool, ...]
output_final_state	bool	是否输出最终状态，默认 False
cu_seqlens	[num_seqs + 1] int32/int64	变长序列的累计长度，varlen 模式必填；相邻重复偏移表示合法的零长度序列
use_qk_l2norm_in_kernel	bool	是否在核内对 Q/K 每个 head 做 L2 归一化（epsilon 1e-6），归一化以 float32 累加再回到输入 dtype；默认 False
output	可选 [total_seq_len, num_o_heads, head_size]	预分配输出，None 时自动分配
output_state	可选 [num_seqs, num_sab_heads, head_size, head_size]	预分配输出状态；output_final_state=True 时必需。给出 state_indices 时是输出状态池，且必须由调用方提供（禁止自动分配紧凑缓冲，否则会被池槽位 id 越界索引）
state_checkpoints	可选 [total_checkpoints, num_sab_heads, head_size, head_size]	中间状态检查点缓冲；checkpoint_every_n_tokens > 0 时必需
checkpoint_cu_starts	可选 [num_seqs + 1] int32/int64	每条序列的累计检查点数，checkpoint_cu_starts[i+1] - checkpoint_cu_starts[i] = seq_len_i // checkpoint_every_n_tokens；要求单调且与 cu_seqlens 一致（该前置条件不在 launch 时校验，以避免设备到主机的同步）
checkpoint_every_n_tokens	int	每 N 个 token 保存一次中间状态，必须是 chunk size（64）的倍数；0 表示禁用检查点（默认）
use_cp	"auto" / bool	是否使用上下文并行；"auto" 为保守路由，True 强制要求 CP 支持，False 禁用 CP
state_indices	可选 [num_seqs] int32/int64	状态池索引（SM90/SM100/SM103/SM120）：序列 i 从 initial_state[state_indices[i]] 读初态并写回同一行（output_state is initial_state 时原地）。id 必须唯一，否则多序列并发写同一池行会留下非确定性的最终状态——唯一性是调用方前置条件，不在 launch 时校验
_cp_chunk_len	可选 int	内部 CP chunk 长度覆盖（测试/调优用），必须是 64 的倍数；None 由后端自动选择
backend	"auto" / "flashinfer" / "cake_gdn" / "cudnn"	后端选择；auto 与 flashinfer 使用同一套 SM90/SM100/SM120 内核与 CP 路由；Cake 内核需显式 cake_gdn；cudnn 走 cuDNN 的融合 SM100 线性注意力引擎（经 flashinfer/cudnn/cudnn_chunk_gated_delta_rule）
max_seqlen	可选 int	最大逻辑序列长度的安全上界，不大于 total_seq_len；CP 内核用它限定每序列 launch grid 而无需从 GPU 读回 cu_seqlens。省略时 CP 使用 total_seq_len（对任意 batch 都正确），传入精确最大值可让 CP 启动更小的 grid

返回值：output_final_state=False 时返回输出张量 [total_seq_len, num_o_heads, head_size]；否则返回元组 (output, final_state)，其中 final_state 在给出 state_indices 时即状态池本身（output_state），其被索引的行现在保存更新后的最终状态。

基础用法：为 decode 准备 recurrent state

一次典型的 GDN prefill 调用（varlen 模式）如下：

import torch
import flashinfer

total_seq_len = 4096
num_seqs = 2
num_q_heads = 16
num_k_heads = num_v_heads = 16
head_size = 128

q = torch.randn(total_seq_len, num_q_heads, head_size, dtype=torch.float16, device="cuda")
k = torch.randn(total_seq_len, num_k_heads, head_size, dtype=torch.float16, device="cuda")
v = torch.randn(total_seq_len, num_v_heads, head_size, dtype=torch.float16, device="cuda")
cu_seqlens = torch.tensor([0, 2048, 4096], dtype=torch.int32, device="cuda")

output, final_state = flashinfer.gdn_prefill.chunk_gated_delta_rule(
    q, k, v,
    cu_seqlens=cu_seqlens,
    output_final_state=True,
)
# output:  [total_seq_len, num_o_heads, head_size]
# final_state: [num_seqs, num_sab_heads, head_size, head_size]

要点：

• cu_seqlens 必须与 q 同设备，支持 int32/int64（调度器会按需把 int64 转成 int32 传给内核，见 flashinfer/gdn_prefill.py）。
• g/beta 缺省时调度器用全 1 的 float32 张量补齐（[total_seq_len, num_sab_heads]）。
• 该内核要求 SM90（Hopper）或 SM100（Blackwell）及以上架构；SM100 路径要求 head_size == 128。

状态池：用 state_indices 接入页式状态管理

state_indices 让调用方无需在「紧凑序列顺序缓冲」和「页式/索引状态池」之间来回 gather/scatter。序列 i 的初始状态读取自 initial_state[state_indices[i]]，最终状态写回 output_state[state_indices[i]]（output_state is initial_state 时即为原地更新）：

N_pool = 64  # 池容量
state_pool = torch.zeros(N_pool, num_sab_heads, head_size, head_size,
                         dtype=torch.float32, device="cuda")
state_indices = torch.tensor([7, 42], dtype=torch.int32, device="cuda")  # 每条序列一个槽位

output, updated_pool = flashinfer.gdn_prefill.chunk_gated_delta_rule(
    q, k, v,
    cu_seqlens=cu_seqlens,
    initial_state=state_pool,
    output_state=state_pool,          # 原地更新
    state_indices=state_indices,
    output_final_state=True,
)

使用约束（来自 API docstring 与 flashinfer/gdn_prefill.py 的校验逻辑）：

• state_indices 仅在 SM90/SM100/SM103/SM120 上受支持，其他架构直接报 NotImplementedError。
• id 必须唯一——这是调用方槽位分配器的责任，内核不做 host 同步校验。
• output_final_state=True 时必须显式传入 output_state 池；调度器拒绝自动分配紧凑 [num_seqs, ...] 缓冲，因为任意池槽位 id 会使其越界。
• SM100/SM103 上，池本身可以是非连续的：任意正的、非重叠的 rank-4 stride 布局都被接受；state_indices=None 时也允许物理视图非连续（见 flashinfer/gdn_prefill.py 与测试 tests/gdn/test_cake_gdn_cp_backend.py 中的 as_strided 用例）。

中间状态检查点

长序列推理（如投机解码、回滚重算）需要在 prefill 过程中按固定间隔保存中间状态。相关三个参数必须成套使用：

checkpoint_every_n_tokens = 256          # 必须是 64 的倍数
seq_lens = [2048, 2048]
total_checkpoints = sum(sl // checkpoint_every_n_tokens for sl in seq_lens)  # 8 + 8 = 16

state_checkpoints = torch.empty(
    total_checkpoints, num_sab_heads, head_size, head_size,
    dtype=torch.float32, device="cuda")
checkpoint_cu_starts = torch.tensor([0, 8, 16], dtype=torch.int32, device="cuda")

output, final_state = flashinfer.gdn_prefill.chunk_gated_delta_rule(
    q, k, v,
    cu_seqlens=cu_seqlens,
    checkpoint_every_n_tokens=checkpoint_every_n_tokens,
    state_checkpoints=state_checkpoints,
    checkpoint_cu_starts=checkpoint_cu_starts,
    output_final_state=True,
)

校验与语义要点：

• checkpoint_every_n_tokens 必须为 64 的倍数，否则直接抛 ValueError；state_checkpoints 与 checkpoint_cu_starts 必须同时提供或同时为 None（flashinfer/gdn_prefill.py）。
• state_checkpoints 形状要求为 [sum(seq_len_i // checkpoint_every_n_tokens), num_sab_heads, head_size, head_size]，dtype 支持 float32/bfloat16/float16/float8_e4m3fn/float8_e5m2。
• 非 Cake 路径下，上下文并行检查点目前支持 SM90、SM100、SM120。
• 调度器会先按 cu_seqlens 与检查点间隔推导期望的 checkpoint_cu_starts 序列，再与传入值比对，不一致即报错（见 flashinfer/gdn_prefill.py）。
• Cake 后端会把「源 chunk 长度 / 物理 CP chunk 长度」与检查点边界映射到固定的 state indices；零长度序列在映射时被跳过（见 tests/gdn/test_cake_gdn_cp_backend.py）。

后端选择：auto / flashinfer / cake_gdn / cudnn

backend 参数决定走哪条实现路径：

• backend="auto"：与 "flashinfer" 完全一致，使用同一套 SM90/SM100/SM120 CuTe-DSL 内核与上下文并行路由，并按低并行度启发式自动决定是否启用 CP（use_cp="auto" 时）。
• backend="cake_gdn"：显式请求 Cake 后端（编译型专用内核）。非 CP 场景下，它要求 Q/K/V 与输出均为连续、rank-3、<a href="https://link.gitcode.com/i/44cd4a667a90ebc26c4520565a971e6c" target="_blank">tokens, heads, 128] 且 FP16/BF16，状态 dtype 必须在 initial/final/checkpoint 间统一（见 [flashinfer/gdn_prefill.py 的 _run_cake_gdn_prefill 校验）。
• backend="cudnn"：委托给 flashinfer/cudnn/cudnn_chunk_gated_delta_rule，即 cuDNN 的融合 SM100 线性注意力引擎；该路径不支持 use_cp=True、checkpoint_every_n_tokens > 0 与 state_indices（flashinfer/gdn_prefill.py）。

值得注意的调度原则（fail-closed）：

• 显式 backend="cake_gdn" 且 use_cp=True 时，若 SM100/SM103 支持条件不满足，直接抛 ValueError，不会静默回退到其他 CP 实现（flashinfer/gdn_prefill.py）。
• use_cp=True 但 CP 不可用（如架构不支持、CUDA 版本不足、内核缺失）时抛 ValueError；而 use_cp="auto" 命中启发式但 CP 不可用时仅发 RuntimeWarning 并回退到非 CP 路径（flashinfer/gdn_prefill.py）。
• 非 Cake 路径在 SM100 上要求 CUDA 13+，否则抛 NotImplementedError（flashinfer/gdn_prefill.py）。

Cake GDN CP 后端：SM100a/SM103a 上的上下文并行

启用方式与路由规则

在 SM100a/SM103a 上显式启用 Cake 上下文并行：

output, final_state = flashinfer.gdn_prefill.chunk_gated_delta_rule(
    q, k, v,
    cu_seqlens=cu_seqlens,
    output_final_state=True,
    backend="cake_gdn",
    use_cp=True,
)

路由规则（docs/api/gdn_prefill.rst）：

• backend="cake_gdn", use_cp=True → Cake GDN CP。
• backend="auto" / backend="flashinfer" → 保留既有 FlashInfer CP 路由（SM90/SM100/SM120 的 CuTe-DSL CP 实现）。
• backend="cake_gdn", use_cp=False 或 "auto" → Cake 非 CP 后端。

调度器判定 use_gdn_cp_backend 需要同时满足：backend 为 cake_gdn、CP 生效、设备计算能力为 (10, 0) 或 (10, 3)、状态/检查点配置落入 FP32/FP16/BF16 域、形状满足 CP 契约（flashinfer/gdn_prefill.py）。

覆盖的公共输入域

Cake GDN CP 后端覆盖 SM100a/SM103a 上完整的合法公共上下文并行输入域，包括：

• 公共 dtype 组合：FP16/BF16 输入，FP32/FP16/BF16 状态，FP32 检查点；
• 等头（equal-head）、GQA、GVA 等 head 映射；
• 可选 gate（g/beta）与 scale；
• int32/int64 的 cu_seqlens 与状态索引；
• 混合（含零长度）乃至全空的 varlen batch；
• 核内 Q/K L2 归一化；
• 检查点、输出；
• 任意正的、非重叠的 packed/indexed/in-place 状态 stride。

对「合法」形状的判定在 flashinfer/gdn_prefill.py 的 _is_gdn_cp_sm100_supported_shape：head_size 必须为 128，head 数满足 (hq == hk and hv % hq == 0) 或 (hk == hv and hq % hk == 0)。测试 tests/gdn/test_cake_gdn_cp_backend.py 用大量 (hq, hk, hv) 与序列长度组合验证了 plan 边界的与 payload 无关性。

四阶段执行流程

Cake GDN CP 保留了如下流水线（manifest 只记录 source inventory 与加载元数据，不携带运行时逻辑）：

T precompute → MN precompute → state fixup → CP prefill

• T precompute：按 64-token chunk 计算 T（key 相关量）；
• MN precompute：在 CP chunk 边界上预计算 M/N；
• state fixup：将各 CP chunk 的递推边界状态修正为一致（含 state_fixup_simt_row4 等 SIMT fixup 内核）；
• CP prefill：真正的上下文并行 prefill 主体（如 cp_prefill_equal_head_checkpoint、cp_prefill_generic_checkpoint、cp_prefill_bf16、cp_prefill_generic_bf16 等按 head/dtype/tail 选择的变体）。

plan 的构建在 flashinfer/jit/cake_gdn_cp_backend.py 的 KERNELS 注册表与 flashinfer/gdn_kernels/blackwell 的 cake_gdn_cp_backend 中完成；tests/gdn/test_cake_gdn_cp_backend.py 验证了各 seq_lens/dtype/head 组合选出的 T/MN/prefill 内核名与 num_sab_heads。

架构与 CUDA 版本要求

• SM100a/SM103a 上，gdn_cp 路由使用 checked-in 的 CUDA 源码（位于 csrc/gdn/gdn_cp，每个 cake_gdn_cp_<hash>_kernel.cu 与 _binding.cu 成对交付），并在 CUDA 12.8、CUDA 12.9、CUDA 13 下受支持，对应 FP16/BF16 输入 + FP32/FP16/BF16 状态 + FP32 检查点。
• 其他 SM100 上下文并行 DSL 路由（包括 FP8 状态或检查点）需要 CUDA 13 与 nvidia-cutlass-dsl[cu13]>=4.4.2（pip install flashinfer-python[cu13]）。
• JIT 侧，每个生成程序按架构编译，架构 flag 集是 JIT 库名的一部分，因此两个架构永远不会共享同一个缓存库（flashinfer/jit/cake_gdn_cp_backend.py）。测试 tests/gdn/test_cake_gdn_cp_backend.py 断言注册表中的每个逻辑内核都能映射到 csrc/gdn/gdn_cp 下真实存在的源文件。

内部机制：plan 缓存、地址重绑定与 CUDA Graph

文档明确了后端的关键运行时设计：

• plan/workspace 缓存 + 直接启动：公共 dispatcher 缓存形状相关的 plan 与 workspace，但直接启动原生 composite。因此 API 分配的输出、以及轮换（rotating）的输入缓冲可以改变地址而无需重建 plan。
• CUDA Graph replay：固定地址的内部 prepared 对象可以使用 CUDA Graph 回放；测试 tests/gdn/test_cake_gdn_cp_backend.py 验证了捕获期（capture）内使用 torch.inference_mode() 产生的元数据张量也能命中缓存，且只在首次调用时 prepare 一次、后续调用直接 replay。
• 索引原地状态走直接 composite：indexed in-place 状态始终留在直接 composite 路径上，因此 preparation 不会推进别名化的 recurrent state。
• TMA 描述符经 __grid_constant__ 传递：每个 TMA 支撑的阶段都把其 descriptor 通过 CUDA 的 __grid_constant__ 内核参数 ABI 传入；后端不保留进程生命周期的描述符 arena，避免描述符泄漏与生命周期管理负担。
• 主机端元数据解析：Cake 后端在 CUDA Graph 捕获之外把 cu_seqlens 等不可变整数元数据解析到主机（LRU 上限 1024 条，见 flashinfer/gdn_prefill.py）；捕获期间若首次遇到未解析的元数据则抛 CakeGDNUnsupportedError，要求调用方在捕获前先做一次 eager 解析。

约束与注意点

• 架构门槛：GDN prefill 仅实现于 SM90、SM100、SM103、SM120；其他架构抛 NotImplementedError。
• head_size：SM100 路径强制 head_size == 128；Cake 非 CP 路径同样要求精确的 [tokens, heads, 128] 形状。
• dtype 一致性：Cake 路径要求 Q/K/V/输出 dtype 完全一致（FP16 或 BF16），门为 FP32，状态 dtype 在 initial/final/checkpoint 之间统一。
• 连续性与设备：Q/K/V/输出必须连续且位于同一 CUDA 设备；cu_seqlens 必须为一维整数连续张量。
• max_seqlen 合法性：不能小于 ceil(total_seq_len / num_seqs)，也不能大于 total_seq_len（flashinfer/gdn_prefill.py）。
• state_indices 唯一性：多序列共享同一池槽位会并发写同一行，产生非确定性结果；由调用方槽位分配器保证。
• 显式 Cake CP 不回退：backend="cake_gdn", use_cp=True 在不支持的输入/架构上失败关闭（fail closed），绝不静默改走其他 CP 实现。

验证体系

仓库在 tests/gdn 下提供了完整的 GDN 测试矩阵，与本文主题直接相关的包括：

• test_cake_gdn_cp_backend.py：注册表完整性、plan 选择、检查点边界映射、状态布局校验、CUDA Graph 捕获/回放、prepared launcher 的地址重绑定；
• test_prefill_cp_delta_rule.py：上下文并行 delta-rule prefill 的正确性；
• test_prefill_state_indices.py：状态池索引路径；
• test_prefill_delta_rule.py 与 test_prefill_auto_backend.py：非 CP 正确性与自动后端路由；
• test_prefill_qk_normalization.py：Q/K 归一化；
• reference_delta_rule.py：blockwise/delta-rule 参考实现，用于数值对照。

这些测试通过 gdn_prefill 模块 的公共入口与内部调度逻辑联动，是理解各参数语义与后端路由行为最直接的旁证。