GDN 与 Chunked Prefill：为什么 prepare_chunk_indices 会出现在 trace 里

核心结论

prepare_chunk_indices 不是普通 attention 的必经路径。它出现在 Qwen3Next 这类带 GDN / Gated DeltaNet 层的模型里，用来给 变长 batch 的 GDN prefill kernel 准备 chunk 映射表。

一句话理解：

GDN 用 recurrent state 表示历史
  -> prefill 不能逐 token 串行递推
  -> kernel 内部按 64-token 小块并行化
  -> 变长 batch 需要一张 “flat chunk -> sequence/chunk id” 映射表
  -> prepare_chunk_indices 就是在准备这张表

完整图解：

GDN 是什么

普通 causal attention 会让当前 token 的 query 访问历史所有 token 的 key/value。GDN 换了一个表达方式：它维护一个 recurrent state，把历史信息压进状态里。

简化模型可以写成：

state_t = gate_t * state_{t-1} + update(k_t, v_t)
output_t = read(q_t, state_t)

这样做的动机是降低长上下文中 attention/KV cache 的压力。代价是：state 有前后依赖。如果 prefill 按 token 递推，GPU 会被迫走很长的串行链。

这里的 state 不是“每个 token 互相独立的一份状态”。更准确地说，它是 per sequence / per layer / per head 的递推状态。处理第 t 个 token 后会得到一个概念上的 state_t：

state_0 -> state_1 -> state_2 -> ... -> state_N

state_t 表示这一层读到 prefix 1..t 后的历史摘要，所以 state_{t+1} 是在 state_t 基础上继续更新出来的。prefill 过程中需要每个 token 对应的输出；但 prefill 结束后，decode 通常只需要保存这一层的最终 recurrent state。

update(k_t, v_t) 可以先理解成“当前 token 对 state 的写操作”。粗略类比是：

k_t: 写入 state 的地址/方向
v_t: 写入的内容
update(k_t, v_t): 当前 token 写入 state 的增量

最简化的写法可以想成外积：

update(k_t, v_t) ~= k_t outer v_t

真实 GDN/DeltaNet 会更复杂，例如还会有 delta rule、门控、归一化或重计算路径；但从依赖关系上，update_t 就是第 t 个 token 已经算出来、准备写入 recurrent state 的增量。

为什么要切成 64-token chunk

GDN prefill 的关键优化是把逐 token 递推改写成块内计算：

原始链路：token 1 -> token 2 -> token 3 -> ... -> token N
块化链路：chunk 0 -> chunk 1 -> chunk 2 -> ... -> chunk N/64

每个 chunk 内部的 64 个 token 通过 cumsum、triangular solve、recompute 等 kernel 处理，尽量转成 GPU 友好的块内并行计算。块和块之间仍然传递 final state，但依赖链从 N 个 token 缩短到约 N/64 个 chunk。

这里的 64 是 GDN kernel 内部 tile，不是服务调度器的 prefill chunk size。

64-token chunk 里的计算形态变化

可以把朴素 GDN prefill 想成 64 次连续的小更新：

S_t = g_t * S_{t-1} + update(k_t, v_t)
y_t = read(q_t, S_t)

这里的麻烦不是单步公式复杂，而是 S_t 依赖 S_{t-1}。如果按 token 做，token 17 的 state 没出来，token 18 就不能算。

chunk kernel 做的事，是把 64 个 token 的 Q/K/V/G 打包成 64 x d 的 tile，再把块内 causal/decay 关系表达成 64 x 64 的三角结构。直观上，它把：

64 次串行小更新

改写成：

几个围绕 64-token tile 的矩阵/三角矩阵/批量 kernel

所以这不是总计算量从 64 变成 1。更准确地说，总 FLOPs 仍然存在，但 token 维度被展开成矩阵维度，GPU 可以用更大的 kernel 粒度、更连续的访存和更高的并行度处理块内计算。块和块之间仍然要传 final state。

递推公式如何展开成块内计算

先只看一个最小递推：

S_t = g_t * S_{t-1} + u_t

其中 g_t 和 u_t 都来自当前 layer 输入 hidden states 的投影结果。进入 GDN recurrent kernel 之前，q/k/v/g 已经算好，u_t = update(k_t, v_t) 也可以看作已知；未知的是 S_1..S_64。

以 4 个 token 为例：

S1 = g1*S0 + u1

S2 = g2*S1 + u2
   = g2*g1*S0 + g2*u1 + u2

S3 = g3*S2 + u3
   = g3*g2*g1*S0 + g3*g2*u1 + g3*u2 + u3

S4 = g4*S3 + u4
   = g4*g3*g2*g1*S0
   + g4*g3*g2*u1
   + g4*g3*u2
   + g4*u3
   + u4

展开后，S4 不再需要先等 S3 的运行时结果；它可以直接由已知的 S0、g1..g4 和 u1..u4 表达出来。把 u 的影响写成表，就是一个下三角 causal/decay 结构：

       u1        u2        u3        u4
S1     1
S2     g2        1
S3     g3*g2     g3        1
S4     g4*g3*g2  g4*g3     g4        1

扩展到 64 个 token，就是一张 64 x 64 的块内关系表。真实 kernel 不一定显式构造完整矩阵，它可能用 scan、triangular solve、recompute 等方式隐式完成；但思想是一样的：

把“后一个 state 等前一个 state”
改写成
“每个 state 由已知输入和块内 causal/decay 权重直接表达”

所以块内并行不是让依赖消失，而是把依赖从运行时逐 token 等待，变成 GPU 更容易批量处理的结构化计算。

chunked prefill 有两层含义

外层是 scheduler 的 chunked prefill。它把长 prompt 拆成多次 model_forward：

forward step 1: 本轮调度 8000 tokens
forward step 2: 本轮调度 8192 tokens
forward step 3: 本轮调度 4000 tokens

这个长度会变，因为它受剩余 prompt 长度、batch 拼接、max_num_batched_tokens、prefix cache 命中和尾块影响。

内层是 GDN kernel 的 chunk。一个外层 model_forward 进入 GDN 层后，每条 sequence 又按 64 token 切成 kernel chunk：

seq0: 130 tokens -> 3 个 GDN chunk
seq1: 20 tokens  -> 1 个 GDN chunk
seq2: 300 tokens -> 5 个 GDN chunk

GPU kernel 看到的是扁平 chunk 列表，所以需要知道：

flat chunk 0 -> seq0, chunk0
flat chunk 1 -> seq0, chunk1
flat chunk 2 -> seq0, chunk2
flat chunk 3 -> seq1, chunk0
flat chunk 4 -> seq2, chunk0
...

这就是 prepare_chunk_indices 的角色。

块内 state 和块间 final state

一个 64-token chunk 内部，逻辑上仍然有每个 token 对应的 state：

S1, S2, S3, ..., S64

每个 token 的输出也要读对应的 state：

y1 = read(q1, S1)
y2 = read(q2, S2)
...
y64 = read(q64, S64)

但传给下一个 chunk 的，只需要这个 chunk 的最后 state：

chunk0: 计算 y1..y64，并把 S64 传给 chunk1
chunk1: 以 S64 为初始 state，计算 S65..S128

因此“块间传 final state”的意思不是块内只算一个 state，而是跨 chunk 边界只需要传递最后的 recurrent state。

Prefill 和 decode 的差异

GDN 对 decode 更友好，prefill 则需要额外的并行化技巧。

prefill 要把整段 prompt 压进 recurrent state：

S1 -> S2 -> S3 -> ... -> SN

如果逐 token 递推，长 prompt 会形成很长的串行链，所以需要 64-token chunk、scan、triangular kernel 等方式降低块内串行度。

decode 每步通常只新增一个 token：

S_new = g_new * S_old + update(k_new, v_new)
y_new = read(q_new, S_new)

这天然就是一步 state 更新，不需要把 64 个 token 展开成块内矩阵问题。相比普通 attention decode 需要读越来越长的历史 KV cache，GDN decode 主要读写固定大小的 recurrent state。

但这不等于整个模型 decode 没有瓶颈。Qwen3Next 是 hybrid 模型，仍有 full attention 层要读 KV cache；同时 decode 还包括 q/k/v/g 投影、MLP/MoE、通信、采样和调度。更准确的说法是：

GDN 用 prefill 的复杂性，换 decode 的长上下文友好性。

为什么 trace 里它可能看起来很贵

vLLM/FLA 的实现会从 cu_seqlens 计算每条 sequence 需要多少个 64-token chunk，然后调用 .tolist() 交给 Python 构造索引表。

如果 cu_seqlens 是 CUDA tensor，那么：

CUDA tensor -> .tolist() -> Python list

这个动作必须把 GPU 上的长度元信息读回 CPU。数据量通常很小，但它会打断 GPU 异步流水，表现为小 DtoH copy、cudaMemcpyAsync runtime 等待和 cudaEventSynchronize/EventSync。

因此 trace 里 prepare_chunk_indices 的 CPU self 高，不一定表示 Python 算术本身慢，更可能表示：

Python 在等待 GPU 把 metadata 同步回来

什么时候容易触发

这个问题通常需要几个条件叠加：

条件	为什么重要
Qwen3Next/GDN 模型	普通 full attention 不走这套 GDN chunk metadata 路径
prefill / chunked prefill	每个 model_forward chunk 都要重新准备本轮变长 metadata
变长 batch	需要 cu_seqlens 描述每条 sequence 起止位置
cu_seqlens 在 GPU 上	.tolist() 会触发 GPU -> CPU 读回
batch shape 动态变化	@tensor_cache 很难跨 step 命中

所以这不是“chunked prefill 一定慢”，也不是“GDN 一定慢”。更准确的判断是：GDN varlen prefill 的 metadata 构造路径，在动态 batch + CUDA cu_seqlens 场景下容易变成同步点。

和现有页面的关系

• Chunked Prefill 深入分析：解释调度层 chunk size、token budget、attention shape 和 P/D 分离边界。
• 批处理与调度：解释 continuous batching 和调度器为什么需要把长 prefill 切开。
• Token Flow 与 Hidden State：解释 hidden state、Attention、MoE、LM Head 和 Sampling 的主流程。
• DeepSeek MLA：对比另一类降低 KV cache 压力的注意力结构，MLA 保留 softmax attention 语义，而 GDN 更偏 recurrent state。