Fast-AR 码本解码与生成主循环

本章覆盖引擎最核心的算法：fast_decode() 与 generate()。前者是一个无 KV cache 的小 Transformer，后者把 Slow-AR、Fast-AR、采样器三者编排成双 AR 主循环。

1. Fast-AR 的设计约束

先看入口的三道校验（:1103-1116）：模型必须带 fast decoder；输入 hidden 必须恰好是 2560 维；前缀长度必须小于码本数 10。最后一条解释了为什么没有 KV cache：

  • Fast-AR 的序列长度始终 ≤ 10+1 = 11，与 fast_context_length 一致；
  • 每次调用窗口长度 = 前缀码数 + 1，随码本序从 1 增长到 10，形状每次都不同；
  • 对这么短的序列，全量重算 4 层比维护一套 cache 更简单、也不慢。

关键维度：fast_dim=2560、32/8 头、head_dim=128（显式给定，不是 2560/32=80）、attn_scale=1/√128（:1118-1127）。

2. 输入拼装：hidden 投影 + 前缀嵌入 

hidden0 = (2560, 1) 的输入张量
projected = fast_project_in × hidden0       若模型带投影；否则直接用 hidden（:1141-1143）
prefix_emb = get_rows(fast_embeddings, prefix_ids)
x = concat(projected, prefix_emb, dim=1)    沿序列维拼接（:1156）

这就是 Fast-AR 的完整输入约定：

  • 第 0 个位置永远是 Slow-AR 的 hidden（经可选投影到 fast_dim），保证每个码本的预测都以同一个语义条件向量为锚；
  • 第 1..k 个位置是本帧已经生成的码，查一张独立的 fast_embeddings 表；
  • 解码第 cb_idx 个码本时，prefix_tokens = [sem_code, residual_1, …, residual_{cb_idx-1}]（由调用方 generate() 收集），形成严格的”从左到右”依赖链。

位置从 0 开始独立编号（:1159-1161），与 Slow-AR 的全局位置无关。

3. 4 层小 Transformer（:1163-1215）

层结构与 Slow-AR 几乎同构：RMSNorm → 融合 wqkv → 视图切分 → 可选 QK-norm → RoPE → GQA 扩展（4×）→ 缩放点积注意力 → wo → SwiGLU FFN。区别只有两处，但都是本质性的：

  1. 无 KV cache 读写：算完 k/v 后直接对当前窗口做注意力，没有 cpy 进持久槽位那一步（对比 s2_model.cpp:991-1005）；
  2. 因果掩码偏移恒为 0：ggml_diag_mask_inf(KQs, 0)（:1197），因为窗口每次都从头开始，第 i 个位置只能看前 i+1 个位置。

RoPE 的上下文参数使用 fast_context_length=11（:1183-1188）。

4. 输出：仅取最后位置 

fast_last = last_token_view(fast_out) 副本    (fast_dim, 1)
logits = fast_output × fast_last             (codebook_size, 1) = 4096

见 :1217-1223。与 Slow-AR 不同，Fast-AR 的输出矩阵 fast_output 不 tied——它是一张独立的 (4096, fast_dim) 权重。每次只预测”下一个码本的一个码”，所以只取窗口最后位置；取回的 4096 个 logits 交给采样器（:1249-1250）。图通过独立的 fast_sched_（容量 16384）分配计算，同样复用 8 MB 上下文缓冲。

5. generate()：双 AR 主循环全景

generate() 是无状态的编排函数。开始时构造全函数复用的 sem_mask（:29-35）：

sem_mask[:] = -∞
sem_mask[sem_begin..sem_end] = 0     只允许 4096 个语义 id
sem_mask[im_end_id] = 0             额外允许结束标记

它的作用是把 Slow-AR 的 155776 维 logits 限制为”要么说一个语义码、要么结束”，杜绝模型吐回普通文本 token。PromptTensor 的转置（:39-44）和 prefill 已在前两章讲过；下面是主循环结构。

5.1 循环骨架 

main_token = 采样(prefill logits + sem_mask)          （:74）
while main_token != im_end 且 step < max_new_tokens:
  ① RAS 重复检测 → 必要时重采 main_token              （:99-116）
  ② sem_code = main_token - sem_begin（钳到 4095）     （:123-125）
  ③ codebooks_cb = [sem_code]
     for cb_idx = 1..num_cb-1:                        （:131-143）
       fast_decode(state.hidden, codebooks_cb) → 采样 → push
  ④ 本帧 num_cb 个码写入 out.codes（按码本行）          （:145-148）
  ⑤ on_frame 回调；返回 false 则中止                   （:150-163）
  ⑥ step_input = [main_token, cb0..cb9] → model.step  （:165-174）
  ⑦ 下一 main_token = 采样(新 logits + sem_mask)       （:181-182）

5.2 两个值得注意的细节

fast_decode 失败不致命：某一码本解码失败时，剩余码本一律填 0 继续走完本帧（:136-139），而不是中断整次合成——0 是一个合法（虽然无声倾向）的码。

main_token 越界保护：理论上 sem_mask 保证 main_token 一定在语义区间，但 :123-125 仍对 sem_code 做了双向钳制，防御性地处理任何意外 id。

5.3 min_tokens_before_end

apply_mask_and_sample（:57-71）在加完 sem_mask 后，可以额外把 im_end 也置 -∞。循环里用 step < min_tokens_before_end 精确控制：前 N 帧强制不能结束，避免超短输出（:181）。

6. RAS：重复感知采样

TTS 自回归生成有个典型故障：模型卡在同一个语义 token 上反复输出（”啊啊啊”式的口吃）。RAS（Repetition-Aware Sampling） 是 s2.cpp 针对这个问题的对策，实现在 :86-121：

ras_window: 最近 10 个语义 token 的滑动窗口
if (main_token 已存在于窗口中 && main_token 是语义 id):
    使用  temperature = 1.0, top_p = 0.9（比默认 0.8/0.8 更发散）
    在当前 logits + sem_mask 上重新采样 main_token

设计要点：

  • 检测在”使用 token 之前”：新采的 main_token 先与历史窗口比对，确认重复才重采，重采结果直接替换；
  • 只对语义 token 生效（:101 的范围判断），不影响 im_end；
  • 提高温度、放宽 top_p 的意图明确：重复说明当前分布过于尖锐，需要注入随机性跳出循环；
  • 窗口维护在 :118-121：每帧 push、超 10 删最旧。

这是一个轻量启发式而非理论保证，但与第 07 章的采样器结合，构成了生成质量的第一道防线。

7. 帧回调：流式生成的钩子 

fcd.codes = codebooks_cb.data();      // 本帧 10 个码
fcd.frame_index / total_frames;
params.on_frame(fcd) == false ⇒ break;

回调在每帧 10 码生成完毕、但尚未 step 之前触发（:150-163）。Pipeline 的流式合成正是利用它：攒够 stride 帧就做一次窗口化解码并吐出音频块（第 10 章）。回调返回 false 用于取消（如 HTTP 客户端断连），此时立即跳出、已生成帧仍保留在结果里。

8. 结果压实

输出缓冲区按 num_cb × max_new_tokens 预分配（:76），但实际帧数通常更少。结尾 :201-213 把它压成 num_cb × n_frames 的紧凑矩阵（逐码本拷贝），布局为 codes[cb * n_frames + t]，与 codec 期望的索引方式一致。

同时函数还输出耗时指标：prefill 毫秒、循环毫秒、平均每帧毫秒（:185-199），供 Pipeline 汇总 RTF（实时率）。

9. 小结

Fast-AR 是一个”hidden 锚定 + 前缀自回归”的无状态小模型，generate() 则用 sem_mask 约束 Slow-AR、用滑动窗口 RAS 防重复、用帧回调支持流式——三者构成完整的双 AR 驱动。所有”选哪个 id”的动作都落在采样器上，第 07 章就进入这个虽小但决定声音随机性的模块。

关键文件

位置 职责
s2_model.cpp:1099 fast_decode 主体
s2_model.cpp:1141 hidden 投影与前缀拼接
s2_model.cpp:1197 窗口因果掩码（偏移 0）
s2_generate.cpp:11 generate 主循环
s2_generate.cpp:29 sem_mask 构造
s2_generate.cpp:86 RAS 重复感知采样
include/s2_generate.h GenerateParams / FrameCallback 定义
This site uses Just the Docs, a documentation theme for Jekyll.