整体架构与 Dual-AR 数据流

本章先不深入任何一个函数，而是把 s2.cpp 运行时的对象关系、张量形状和一帧数据的完整旅程讲清楚。后面所有章节都是这张图的局部放大。

1. 运行时对象关系

核心 C++ 类只有四个，全部位于命名空间 s2：

Pipeline（s2_pipeline.cpp）
 ├── Tokenizer        s2_tokenizer.cpp   Qwen3 BPE，读 tokenizer.json
 ├── SlowARModel      s2_model.cpp       权重、两套调度器、KV cache
 │     ├── hparams_                       结构超参（维度/层数/码本…）
 │     ├── weights_.ctx_w                 只存元数据的 GGUF 上下文（无数据）
 │     ├── sched_                         Slow-AR 的后端调度器（图容量 32768）
 │     ├── fast_sched_                    Fast-AR 的后端调度器（图容量 16384）
 │     ├── memory_k_ / memory_v_          F16 KV cache（预分配）
 │     └── 后端 buffer 向量                GPU/CPU 权重各若干块
 └── AudioCodec       s2_codec.cpp       卷积 codec + VQ 主机缓存
       ├── Impl（pImpl）                   全部 codec 参数与权重张量
       ├── semantic_vq / residual_vq      预归一化的 VQ 码本（CPU 侧缓存）
       └── decode_cache                   按帧数缓存的解码图

外部消费者有三个：CLI（main.cpp）、HTTP 服务（s2_server.cpp）、C 导出层（s2_export_api.cpp）。它们都只跟 Pipeline 打交道，不直接操作模型。

generate()（s2_generate.cpp:11）是夹在 Pipeline 和模型之间的无状态算法函数：输入模型、生成参数和 PromptTensor，驱动双 AR 循环，输出整数码矩阵。

2. 关键维度一览

这些数字来自 GGUF 元数据，缺失时由 s2_model.cpp:258-288 与 s2_codec.cpp:699-703 给出默认值，是理解后续形状的前提：

维度 值 含义
vocab_size 155776 Qwen3 词表大小
semantic_begin/end 151678 / 155773 语义 token 在词表中的区间（共 4096 个）
embedding_length 2560 Slow-AR 隐状态维度
block_count 36 Slow-AR 层数
feed_forward_length 9728 SwiGLU FFN 中间维度
head_count / head_count_kv 32 / 8 GQA：32 查询头，8 KV 头
head_dim 80（2560/32） 无 QK-norm 时；有 q_norm 张量时以其形状为准
RoPE base 1 000 000 Qwen3 长上下文设定
codebook_size 4096 每个码本的条目数
num_codebooks 10 1 个语义码本 + 9 个残差码本
Fast-AR 4 层，窗口 ≤11，head_dim 128 逐帧码本解码器
采样率 / hop 44100 Hz / 512 codec 帧率
samples_per_frame 2048 一帧 = hop × ∏下采样倍数 = 512×4

注意一帧的两个含义：codec 帧 = 2048 个采样点（约 46.4 ms），而生成端的一帧（一个语义步）对应一个 codec 帧。samples_per_frame_ 的计算在 s2_codec.cpp:859-868。

3. 一帧数据的完整旅程

3.1 输入侧：多行 token 矩阵

S2 系列模型的输入不是普通语言模型的一维 token 序列，而是一个 (num_codebooks+1) × T 的矩阵，共 11 行：

  • 第 0 行：语义流——文本 BPE token、特殊标记，以及参考音频段的语义码（码本索引 + semantic_begin）；
  • 第 1～10 行：码本流——每个位置放该码本的离散码（取值 0..4095），仅在”参考音频占位”段有值，其余位置填 0 并由语义掩码屏蔽。

这个矩阵由 build_prompt() 构造（详见第 09 章），在 s2_generate.cpp:39-44 转置成列优先布局后交给 prefill()：

转置前（行优先）：prompt.data[row * cols + pos]
转置后（列优先）：prompt_tm[col * rows + row]  →  flat_tokens[t * codebook_dim + cb]

模型内部统一使用列优先，codebook_dim = num_codebooks + 1 = 11（s2_model.cpp:799）。

3.2 Slow-AR：一个语义 token，外加 10 个码本嵌入

每个位置 t 的输入向量这样构成（s2_model.cpp:941-959）：

x_t = sem_scale · embedding[ semantic_id_t ]                    （第 0 行查主嵌入表）
    + Σ_{cb=1..10} sem_mask[t] · embedding_cb[ code_id_{t,cb} ] （码本行查各自的小嵌入表）

其中：

  • 10 个码本共用一张 codebook_embeddings 表（10×4096 行），码本 cb 的码 id 要加上偏移 cb * codebook_size 再查表（s2_model.cpp:903-921）；
  • sem_scale = 1/√11、attn_scale = 1/√head_dim（s2_model.cpp:890-891）；
  • 在纯文本位置掩码为 0，码本项不贡献，只有语义嵌入——所以文本段的行为与普通语言模型一致。

36 层 Transformer（RMSNorm → GQA 注意力 → SwiGLU FFN）之后，取最后位置隐状态 hidden（2560 维），输出 logits 由 tied embedding 得到：logits = hidden × embeddings.weight^T（s2_model.cpp:1058）。

3.3 交接：一个语义步内的 10 次 Fast-AR

Slow-AR 每走一步，采样出一个语义 token main_token。它同时是本帧第 0 个码本的码：sem_code = main_token - semantic_begin。然后 Fast-AR 接手（s2_generate.cpp:127-143）：

Fast-AR 输入 = [ Slow-AR 的 hidden(2560),                                    ← 经 fast_project_in
                本帧已生成的前缀码 embeddings（第 0 个=sem_code, 第1..k个…） ]
窗口长度 = 已生成码本数 + 1（≤11），4 层小 Transformer 全量重算（无 KV cache）
输出 = codebook_size=4096 个 logits → 采样得到下一个码本的码

重复 9 次，凑齐第 1～9 个残差码本。于是一个语义步产出 10 个码：[sem_code, residual_1, …, residual_9]。

3.4 回灌与前进

一帧 10 码齐了之后，这 11 个值（语义 token + 10 个码本码）作为下一位置的完整输入，调用一次 Slow-AR 的 step()（s2_generate.cpp:165-174），KV cache 保留全部历史，进入下一帧。如此循环直到采到 <|im_end|> 或达到 max_new_tokens。

        ┌─────────────────────── 帧 n ───────────────────────┐
文本/前缀 ──Slow-AR(step)──► main_token ──► Fast-AR×9 ──► 10 码
                                  ▲                            │
                                  └──── 11 值回灌 + KV 增长 ◄──┘

3.5 输出侧：码矩阵 → 波形

生成结果 codes 按码本行存放：codes[cb * n_frames + f]（s2_generate.cpp:145-147）。codec 解码时：

10 × T 离散码
 │  quantizer decode：语义VQ(quantizers.0) + Σ 残差VQ(quantizers.i)，逐级相加 + 上采样
 ▼
 latent（卷积序列）
 │  waveform decoder：ConvNeXt/残差块上采样，tanh 限幅
 ▼
 T × 2048 个 44.1kHz 采样点

解码图的拼装在 build_decode_codes_stage_backend()，完整 decode() 在 s2_codec.cpp:1297。

4. 两种”图”的节奏

ggml 的执行模型是”建图 → 调度 → 计算”。s2.cpp 针对不同阶段用了不同的缓存策略，这是理解性能特性的关键：

路径 图的生命周期 上下文预算 说明
Slow-AR step() 复用（eval_cached） 约 10 MB 可复用 ctx 每帧形状不变（单步 11 行），建一次反复跑
Slow-AR prefill() 每次新建，可分块 随 chunk 大小 CUDA 强制 1 token/块；其他后端自适应 8～64
Fast-AR fast_decode() 每次新建 约 8 MB 窗口长度随码本序增长（1→10），图必变；无 KV cache
Codec GPU 解码 按 n_frames 缓存 256 MB fused 图，失败的帧数拉黑
Codec CPU 解码 两张图分别建 96 + 128 MB 量化器解码、波形解码分离
Codec 编码 两张图 128 + 96 MB 编码器图 + RVQ 前处理图

调度器本身是成员级单例：Slow-AR sched_ 图容量 32768、Fast-AR fast_sched_ 16384（s2_model.cpp:595-602）。

5. 后端与放置策略概览

后端枚举定义在 s2_backend.h：CPU = -1、Vulkan = 0、CUDA = 1、Metal = 2。权重放置有一条容易被忽略但很重要的策略线：

  • CUDA：嵌入表（embeddings / codebook_embeddings / fast_embeddings）与最终 norm 强制留在 CPU（s2_model.cpp:486 起、README.md:45）。原因是 K 量化嵌入的 CUDA get_rows 不被支持，且曾经导致声音克隆 prefill 不稳定；语义 prefill 也因此强制逐 token（s2_model.cpp:37）。
  • Vulkan / Metal：非全卸载时同样保留嵌入；当所有 36 层都上 GPU（full_model_offload）时，整个模型都在 GPU（s2_model.cpp:481）。
  • 当驱动对单块 buffer 有大小上限时，权重会被自动切成多块放置（allocate_weight_buffers()），而不是申请一块巨型连续显存。
  • codec 默认 --codec-auto：用 8 帧零码在 CPU/GPU 上各跑一次，只有 GPU 快于 0.90× CPU 时间才用 GPU；GPU 初始化或分配失败一律退回 CPU（详见第 10 章）。

6. 小结

s2.cpp 的架构可以浓缩成三句话：

  1. 矩阵进、矩阵出：11 行 token 矩阵是 S2 多码本建模的统一接口，文本与参考音频在同一布局下共存；
  2. 一大一小两个 AR：Slow-AR 带 KV cache 跨步前进、每步交接一个隐状态；Fast-AR 无状态地在每步内全量重算、逐码本解码；
  3. 一切皆 ggml 图：生成路径复用图、prefill 分块建图、codec 按帧数缓存图——图的缓存粒度基本决定了各阶段的性能行为。

第 03 章先落地到仓库本身：代码怎么组织、怎么构建、为什么构建时要给 ggml 打补丁。

关键文件

文件 职责
include/s2_model.h SlowARModel 类、hparams、权重结构
include/s2_pipeline.h Pipeline 与 PipelineParams
src/s2_prompt.cpp 11 行 PromptTensor 构造
src/s2_generate.cpp 双 AR 交接的主循环
src/s2_codec.cpp 码 → 波形解码
This site uses Just the Docs, a documentation theme for Jekyll.