音频编解码器:卷积网络、RVQ 与缓存图

s2_codec.cpp 是仓库中最大的文件(1447 行),实现 S2 Pro 的神经音频 codec:44.1 kHz 波形 ↔ 10 个离散码本。它在流水线中承担两个方向的职责——声音克隆时把参考音频编码成码,合成时把生成的码解码回波形。

1. 结构总览

codec 是一个典型的”卷积主干 + Transformer 中段 + 残差向量量化”架构:

编码方向:
  波形 (44.1kHz)
   → initial conv → encoder blocks(卷积下采样 + 可选 window transformer)
   → final conv → 下采样 stages → RVQ pre_module transformer
   → 语义VQ + 9×残差VQ(逐级减去投影)→ 10 码本 × T

解码方向(逆过程):
  10 码本 × T
   → 语义VQ decode + Σ 残差VQ decode(逐级相加)
   → post_module transformer + 反向上采样
   → decoder blocks(转置卷积上采样)→ tanh → 波形 (T×2048)

全部实现参数保存在 pImpl 结构 AudioCodec::Impl,包括编解码维度、各阶段速率、VQ 主机缓存和解码图缓存。

2. 卷积构件

2.1 因果卷积与 Snake 激活

codec 的卷积都是因果的(输出位置 t 只依赖输入 ≤t 的位置),这对低延迟流式场景是必要的:

通道布局在卷积(channels-last)和线性层(channels-first)之间用 cl_to_lc / lc_to_cl 显式转置,配套一组 channel 偏置/缩放/RMSNorm 的 cl 版本辅助函数。

2.2 三个基本块

构件 组成 位置
Residual unit 两层因果卷积 + dilations(1,3,9)跳跃连接 :379
ConvNeXt block 深度可分离 conv → layer_norm(1e-6) → pwconv1 + GELU(erf) → pwconv2 + γ :393
Encoder/Decoder block 卷积(转置)步长 + 残差单元;编码块可挂 window=512 的 transformer :420、:480

完整解码器 build_decoder() 在输出端套 tanh 限幅。

3. Codec 内的 Transformer

量化器前后各有一个小 Transformer,由 build_transformer() 统一构建。层数不靠参数指定,而是循环到 wqkv 张量缺失为止(:301-304),直接以权重文件的实际内容为准。

它与语言模型 Transformer 有两处显著差异:

  1. 每层要求 layer-scale:必须存在 attention_layer_scale.gamma 和 ffn_layer_scale.gamma(:318-319),注意力和 FFN 输出在残差相加前各乘一个逐通道 γ(:361、:369);
  2. FFN 用 SiLU 而非 SwiGLU:silu(gate) ⊙ up 是手工拼的(:367),因为它的权重只有 w1/w2/w3 但 gate 语义不同;没有用融合 swiglu 算子。

3.1 掩码:窗口、全因果、或交给 diag_mask

prepare_transformer_inputs() 决定注意力掩码形态:

window_size > 0 且 < seq_len  → 滑窗掩码(只能看 [q-window+1, q])
否则后端要求显式因果掩码       → 全因果掩码(-1e9 填充)
两者都否                       → 建图时用 diag_mask_inf(:352)

显式掩码目前只对 Metal 构造(backend_requires_explicit_causal_mask()),因为该后端在动态形状下内建因果掩码支持有限;其他后端用零开销的 diag_mask_inf。掩码矩阵一次性在主机侧构造(:278-284),值用 -1e9 而非 -inf。

4. 向量量化:余弦 VQ

codec 的 VQ 不是最邻近欧氏距离,而是余弦匹配 + 投影的形式。主机侧缓存由 load_vq_cache() 建立,除了读入 in/out 投影和原始码本,还预计算一份逐条目 L2 归一化的码本(:557-564),避免每次编码重复归一化。

4.1 quantize_with_vq()

编码单级 VQ 的完整过程(:585-618):

in_proj: input(in_dim) → latent(codebook_dim)
对每帧:
  latent L2 归一化
  score_k = normalized_latent · normalized_codebook[k]   余弦相似度
  code = argmax_k score
  输出投影:out_proj( codebook[code](原始、未归一化向量)) → input_dim

要点:匹配在归一化空间(余弦),但还原用的是码本原始向量——归一化只用于”选谁”,不改变被表示的内容。投影矩阵用朴素的 project_1x1() 在 CPU 上计算,所以 VQ 编码不依赖后端算子支持。

4.2 RVQ:逐级相减

残差量化的核心是每级量化后从输入中减去该级的投影输出,下一级只处理残差(编码主循环 :1136-1154):

residual = stage
(code, proj) = VQ_semantic(residual); residual -= proj     语义码本行 0
for i in 0..8:
  (code, proj) = VQ_residual[i](residual); residual -= proj 码本行 i+1

9 级残差逐级逼近,码本索引布局为 codes[cb*T+t]。这就是”1 语义 + 9 残差 = 10 码本”的来历,num_codebooks_ = residual + 1(:858)。

4.3 解码侧:逐级相加

build_decode_codes_stage_backend() 是编码的镜像:

  • 第 0 路:build_vq_decode_stage() 读语义码本(quantizer.semantic_quantizer.quantizers.0):get_rows(codebook, ids) → out_proj;
  • 第 1~9 路:读 quantizer.quantizer.quantizers.i,每路结果 ggml_add 累加(:655-666);
  • 10 路在同一级相加,等价于把残差重新拼回。随后接 post_module transformer 和反向上采样。

build_vq_decode_stage 还处理了 out_proj 权重以 4D 存储的情况(:632-636),reshape 成 2D 再做线性投影。

5. 加载与元数据

load_shared() 与模型共享同一 GGUF。开头先重置全部状态(:699-703),GPU 初始化失败时回退 CPU(:710-748),并复用模型的 ctx_w 或自行打开 GGUF(:750-762)。

5.1 tprefix:两种架构

if (arch == "fish-speech")        tprefix = "c.";       // 合并文件:codec 张量在 c.* 下
else if (arch == "fish-speech-codec") tprefix = "";     // 独立 codec 文件:无前缀

见 :815-822。后续所有张量名拼接都基于 tprefix——这解释了导出脚本为什么给 codec 张量保留 c. 前缀。

5.2 参数与帧尺寸

:824-854 读取一大组 fish_speech.codec.* 元数据:采样率、hop、编/解码维度、各级速率、量化器维度与码本大小、两组 Transformer 参数。随后计算:

samples_per_code_frame = hop × ∏downsample_factor = 512×4 = 2048   (:859-868)
streaming_history_frames = round_up_8(rvq_window - 1) + 16         (:870-874)

streaming_history_frames_(典型 160)是流式解码的默认上下文长度,第 10 章会用到。权重 buffer 用 ggml_backend_alloc_ctx_tensors 一次性分配到选定后端(:876-877)。

数据读取后,refresh_host_caches() 重建语义/残差 VQ 的主机缓存——每次权重重置都要刷新,因为余弦匹配依赖码本内容。

6. encode():两张图

encode() 先把输入补齐到 frame_length 的整数倍,然后分两张图计算:

  1. 128 MB 编码器图:initial conv → encoder blocks → final conv,产出 latent(:999-1022);
  2. 96 MB 量化器前处理图:下采样 stages → pre_module RVQ transformer(:1080-1095),产出供 VQ 的 stage;
  3. 在主机侧跑 4.2 节的语义 + 残差 VQ 序列,得到全部码。

拆两张图的原因与 prefill 相同:单张 128+96 MB 的大图中间物化过多;而 VQ 是 CPU 主机计算,天然形成图的边界。

7. decode():缓存图优先,split 路径兜底

decode() 有两条路径,这是 codec 性能设计的核心。

7.1 GPU fused 缓存图

非 CPU 后端首先尝试 run_cached_decode_graph():

  • 图由 build_cached_decode_graph() 构建,256 MB 上下文、quantizer decode + waveform decode 融合为一张图;
  • 按 n_frames 缓存:同一帧数(流式场景中高度重复的窗口尺寸)直接复用已建图,不必重建;
  • 失败拉黑:任何阶段失败都记录 failed_n_frames = n_frames(:1198 起),下次同帧数直接跳过、不再重试。重置缓存时拉黑记录会保留(reset_decode_cache :128-131);
  • 解码前先经 sanitize_decode_codes():把超出各码本大小的非法码钳回合法范围,防止坏码触发后端错误(:1262)。

7.2 CPU/split 两张图

fused 图不可用(CPU 后端、或该帧数曾失败)时走 split 路径:

  1. 96 MB 量化器解码图:codes → VQ decode 逐级相加 → post_module transformer,产出 latent(:1322-1331);
  2. 128 MB 波形解码图:latent → decoder blocks → 音频(:1400-1410)。

两条路径输出形状一致,调用方无感知。GPU 路径任何失败都会自然落回 split(CPU)路径,因此即使融合图对某个尺寸不支持,合成也不会中断——这与”GPU codec 失败一律退回 CPU”的总策略一致。

8. 小结

codec 模块可以归纳为三点:因果卷积主干保证流式可行性、余弦 VQ + 逐级残差构成离散音频接口、按帧数缓存的 fused 图与 split 兜底图兼顾性能与鲁棒性。它的输出既直接写成 WAV,也通过帧回调进入流式管线。第 09 章回到输入端,看看参考音频是如何与文本一起组成模型 prompt 的。

关键文件

位置 职责
s2_codec.cpp:288 codec Transformer(layer-scale + SiLU)
s2_codec.cpp:585 余弦 VQ 量化
s2_codec.cpp:646 码本逐级相加(解码侧)
s2_codec.cpp:693 load_shared 与元数据
s2_codec.cpp:980 encode(两图)
s2_codec.cpp:1190 GPU fused 缓存图
s2_codec.cpp:1297 decode 入口与路径选择

This site uses Just the Docs, a documentation theme for Jekyll.