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

prepare_chunk_indices 不是普通 Attention 的通用路径。它服务于 Qwen3Next 这类 GDN / Gated DeltaNet 层：把变长 batch 中的 token 切成 kernel 内部小块，并建立扁平 chunk 到 sequence 的映射。

理解这段 trace 的关键，是先分开服务调度器的 Prefill chunk与 GDN 算子内部的 kernel chunk。

1. 先看两层 chunk

外层由 serving scheduler 决定：

model_forward 1: request A 推进 8K token
model_forward 2: request A 推进 8K token + request B decode
model_forward 3: request A 推进剩余 token

它受 token budget、batch 拼接、prefix hit、KV 水位和尾块影响，详见 Chunked Prefill 深入分析。

内层由 GDN kernel 决定。一个外层 forward 进入 GDN 层后，每条 sequence 再按实现粒度切块：

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

已有实现中常见的 64 token 是 kernel tile，不是 scheduler 的 Prefill chunk size。两者即使都叫 chunk，也不能用同一个配置或 trace 计数解释。

2. GDN 为什么需要块内并行

普通 causal Attention 直接访问历史 K/V。GDN 则维护 recurrent state，把历史压进状态：

$$S_t=g_tS_{t-1}+\operatorname{update}(k_t,v_t), \qquad y_t=\operatorname{read}(q_t,S_t)$$

这里的 state 是 per sequence / layer / head 的递推状态。若 Prefill 逐 token 执行，$S_t$ 必须等待 $S_{t-1}$，长 Prompt 会形成长串行链。

Kernel chunk 将一组 token 的 Q/K/V/G 打包成 tile，用 scan、triangular solve 或 recompute 等结构化计算处理块内依赖。依赖没有消失，只是从运行时逐 token 等待，改写成 GPU 更容易批量执行的形式。

3. 递推如何展开成三角结构

先看最小递推：

$$S_t=g_tS_{t-1}+u_t$$

四个 token 展开后：

$$\begin{aligned} S_1 &= g_1S_0+u_1 \\ S_2 &= g_2g_1S_0+g_2u_1+u_2 \\ S_3 &= g_3g_2g_1S_0+g_3g_2u_1+g_3u_2+u_3 \\ S_4 &= g_4g_3g_2g_1S_0+g_4g_3g_2u_1+g_4g_3u_2+g_4u_3+u_4 \end{aligned}$$

$u$ 对各 state 的影响形成下三角 causal/decay 关系：

$$\begin{array}{c|cccc} & u_1 & u_2 & u_3 & u_4 \\ \hline S_1 & 1 & 0 & 0 & 0 \\ S_2 & g_2 & 1 & 0 & 0 \\ S_3 & g_3g_2 & g_3 & 1 & 0 \\ S_4 & g_4g_3g_2 & g_4g_3 & g_4 & 1 \end{array}$$

扩到 64 token，就是对应的块内结构。真实 kernel 不一定显式构造完整矩阵，但这个展开解释了为什么 token 维可以转成矩阵/scan 计算。

4. 块内 state 与块间 final state

一个 chunk 内仍需要每个 token 的输出：

$$y_t=\operatorname{read}(q_t,S_t)$$

但传给下一个 chunk 的只需当前块的最终状态：

$$(y_1,\ldots,y_{64};S_{64}) \longrightarrow (y_{65},\ldots,y_{128};S_{128})$$

所以“块间只传 final state”不等于块内只计算一个 state；它描述的是跨 chunk 边界的最小依赖。

5. prepare_chunk_indices 的具体职责

变长 batch flatten 后，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 根据 cu_seqlens 等长度元数据，生成这类映射，使 kernel 能找到：

• chunk 属于哪条 sequence；
• 在该 sequence 中的块位置；
• 首尾块是否不完整；
• 应读取哪个 incoming state、写回哪个 final state。

这属于 GDN varlen kernel metadata，不是 scheduler 选择请求的步骤。

6. Prefill 与 Decode 的差异

Prefill 要把整段 Prompt 压进 recurrent state，因此需要块内并行化来缩短串行链。

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

$$S_{\text{new}} =g_{\text{new}}S_{\text{old}} +\operatorname{update}(k_{\text{new}},v_{\text{new}})$$

单个 GDN 层只需读写固定大小的 recurrent state，不必扫描不断增长的 KV。但 Qwen3Next 是 hybrid 模型，仍可能包含 full-attention、MLP/MoE、通信、采样与调度；不能据此推导整模型 Decode 一定不受带宽或 KV 影响。

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

已检查的 vLLM/FLA 路径会从 cu_seqlens 计算每条 sequence 的 chunk 数，并通过 .tolist() 交给 Python 构造索引。

若 cu_seqlens 位于 CUDA：

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

CPU 必须等 GPU 元数据可见。即使拷贝的数据很小，也可能暴露此前排队的 GPU 工作，在 trace 中表现为 DtoH copy、runtime 等待或 event synchronization。

因此：

• prepare_chunk_indices CPU self time 高，不一定代表 Python 算术慢；
• 同步与该调用同时出现，说明它是一个可疑边界；
• 没有时间线、调用栈和对照实验，不能断言它就是端到端损失的唯一根因。

8. 触发条件与验证

条件	意义
GDN 模型	普通 full-attention 不走同一 metadata 路径
Prefill / varlen batch	需要处理多条不同长度 sequence
CUDA cu_seqlens	转 Python list 可能触发读回
动态 shape	缓存难以跨 step 命中

最小验证应比较：

1. 相同 token 总量下，单一长度与变长 batch；
2. metadata 留在 device、提前在 host 维护或原始 .tolist() 路径；
3. CPU self、GPU idle gap、DtoH/event wait 与端到端 TTFT；
4. 相同请求、模型、版本和采集窗口。

只有同步减少且端到端指标一致改善，才可把它升级为已验证因果结论。

相关页面

• Attention 架构演化
• Chunked Prefill 深入分析
• 批处理与调度
• Prefill Trace
• Token Flow 与 Hidden State