Mechanism Diagram

为什么 GDN prefill 要切成 chunk，prepare_chunk_indices 又在准备什么

这张图把三个概念放到一条链路里：Qwen3Next 的 GDN 层、外层 scheduler 的 chunked prefill、GDN kernel 内部的 64-token 小块。

面向 trace 解释：CPU self 高 ≠ 纯 Python 算术慢

1. Qwen3Next 里不是所有层都是普通 Attention

一个模型 forward 会扫过很多层 示意：每 4 层里有 3 个 GDN / linear attention 层，1 个 full attention 层 GDN state GDN state GDN state Full Attn KV cache ... GDN state GDN state GDN state Full Attn KV cache GDN 层在 prefill 时要把序列历史压进 recurrent state
GDN：状态递推层 Full Attention：普通注意力层

2. 逐 token 递推太串行，所以要块内并行

朴素算法：token 级串行 t1 t2 t3 t4 ... 后一个 state 依赖前一个 state，长 prompt 会变成很长的等待链。 GDN prefill：把 64 个 token 变成一个块内矩阵问题 chunk 0 token 0..63 chunk 1 token 64..127 chunk 2 token 128..191 块内尽量并行；块间只传 final state，依赖链从 N token 缩短到 N/64 chunk。

3. 64-token chunk 里，计算形态怎么变

朴素理解：64 次 token 级小更新 每一步读旧 state，写新 state；后一步必须等前一步完成。 token t0 S0 -> S1 token t1 S1 -> S2 token t2 S2 -> S3 ... 每个 token 的简化递推 S_t = g_t * S_{t-1} + update(k_t, v_t) y_t = read(q_t, S_t) 计算量还是 O(64)，但关键问题是 64 个 state 依赖点串在一起。 chunk kernel：把 64 个 token 打包成 tile 把 token 维度显式展开，交给 GPU 做矩阵/三角矩阵/批量 kernel。 Q 64 x d 64 tokens K,V,G 64 x d 批量输入 T 64 x 64 causal / decay Y 64 x d 批量输出 可以近似理解为 64 次小更新 -> 几个更大的 tile kernel / GEMM-like 操作 总 FLOPs 不会凭空消失，但并行度、访存连续性和 kernel 粒度变好。 直观规模对比 串行版：64 个依赖边，GPU 很难同时算 token 1..63 的 state。 块化版：块内 token 维度变成矩阵维度，块间只剩 final state 依赖。

这里的“矩阵乘法”是直观近似：真实 FLA/GDN kernel 会包含 cumsum、triangular solve、recompute 等步骤，不是把全部递推简单替换成一个 GEMM。

4. 两层 chunk：外层 scheduler 动态切 prefill，内层 GDN 固定按 64 token 切 kernel chunk

外层：vLLM scheduler 每次凑一个 model_forward 每次 token 数可能变：受剩余 prompt、batch 拼接、prefix hit、尾块、max_num_batched_tokens 影响。 model_forward step 1 req A 5000 + req B 3000 = 8000 tokens model_forward step 2 req A 剩余 8192 tokens model_forward step 3 req A 2800 + req C 1200 = 4000 tokens 内层：一个 step 进入 GDN 层后，再按 64-token kernel chunk 切每条 sequence 这里的 64-token chunk 是 GDN kernel 的 tile，不是外层 scheduler 的 prefill chunk。 seq0: 130 tokens 0..63 64..127 tail seq1: 20 tokens tail seq2: 300 tokens 0..63 64..127 128..191 192..255 tail prepare_chunk_indices 输出的映射表 kernel 只看到扁平 chunk，需要知道每个 chunk 属于哪条 seq。 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 ...

5. 为什么这个实现会在 trace 里表现成 CPU self 高

cu_seqlens 在 GPU 上 每条 seq 的起止位置 GPU tensor 上算 num_chunks ceil(length / 64) .tolist() Python list 必须在 CPU Python for n in list GPU → CPU 读回 + 同步等待 你的 trace 里小 DtoH 拷贝数据量很小，但 cudaMemcpyAsync / EventSync 等待很长，所以 CPU self 主要是同步等待。
GDN 的设计目标

用 recurrent state 表示历史，降低长上下文 attention 的计算/访存压力。

为什么要 64-token chunk

把逐 token 递推改写成块内并行计算，块间只传 final state。

prepare_chunk_indices 的角色

给变长 batch 的 GDN kernel 提供 “flat chunk → sequence/chunk id” 映射。