代码解读 | CosyVoice 代码研读（三）：CosyVoice DiT CFM
白御青 Lv5
  2026-01-29 2026-01-29   2026-03-27 12:10:58      7.4k 字  29 分钟  

CosyVoice3 的训练代码已经开源，这次最核心的变化在于 Flow Matching 模块从 U-Net 升级为 DiT（Diffusion Transformer）架构。本文以 cosyvoice/flow 目录下的最新代码为切入点，结合配置文件 cosyvoice3.yaml，从模型架构、Flow Matching 原理、条件信息输入机制、流式推理设计等维度对 DiT 模块进行解读。由于 DiT 的实现复用了 CosyVoice 1 和 CosyVoice2 中的不少关于 CFM 的基础代码，本文也一并进行简要分析。

1. 整体架构概览

CosyVoice 系列的语音合成模型分为三个模块：LLM→Flow Matching→HiFiGAN(HiFTNet)，其中 Flow Matching（后文有可能简称为 CFM） 模块负责将离散的 speech token 转化为连续的梅尔特征，因为 HiFiGAN 声码器技术已经相对成熟且效果不错，所以合成音频的音色、音质基本上可以认为是 Flow Matching 这个模型来保证的。本文基本不涉及 Flow Matching 的基础理论，主要关注 CosyVoice 里的具体代码实现。

在 CosyVoice 开源代码下，Flow Matching 模块的顶层类为 CausalMaskedDiffWithDiT（定义于 cosyvoice/flow/flow.py），其内部结构如下：

1
2
3
4
5
6
CausalMaskedDiffWithDiT
├── input_embedding: nn.Embedding(6561, 80)              # speech token → 80维向量
├── spk_embed_affine_layer: Linear(192 → 80)             # speaker embedding
├── pre_lookahead_layer: PreLookaheadLayer(80→1024→80)   # 因果卷积限制模型感受野
└── decoder: CausalConditionalCFM                        # 条件 Flow Matching
    └── estimator: DiT                                   # DiT 网络（非自回归 Transformer）

与 CosyVoice 前代模型的架构变化

版本 flow 类 编码器 上采样方式 去噪网络
v1 MaskedDiffWithXvec ConformerEncoder + LengthRegulator 线性插值 U-Net
v2 CausalMaskedDiffWithXvec UpsampleConformerEncoder 转置卷积 + Conformer U-Net
v3 CausalMaskedDiffWithDiT PreLookaheadLayer repeat_interleave DiT

可以看出，CosyVoice3 的核心变化是：用轻量 PreLookaheadLayer 替代了 Conformer Encoder，用 DiT 替代了 U-Net，将模型参数集中到 Transformer 网络上。

2. 模型结构解读

本章按照模型从输入到输出的数据流向，依次介绍输入信息、注入模型的方式，token 编码与上采样、DiT 网络结构，训练目标和推理采样方式。

2.1 输入条件与信息融合

DiT 在推理/训练时一共接收五路信息，它们的注入方式和生命周期存在一些差异：

参数名 信息 注入方式 注入位置 随 ODE 迭代变化
mu（token emb） speech token 信息 concat → Linear DiT 输入 不变
spks（speaker） 说话人表征 concat → Linear DiT 输入 不变
cond（prompt mel） 参考 prompt mel concat → Linear DiT 输入 不变
x（噪声/状态） ODE 当前状态 concat → Linear DiT 输入 每步更新
t 时间步 AdaLN-Zero 调制 DiT 每一层 从0到1递增

上述信息进入网络，主要通过：

  • 方式一：concat 投影，mu / spks / cond / x 在入口 InputEmbedding 中拼接为 320 维后线性投影到 1024 维，之后仅通过 attention 间接影响各层。
  • 方式二：AdaLN-Zero 调制，时间步 t 在每一层通过 6 个调制参数（shift/scale/gate × attn/FFN），能够在每个深度上调制影响模型能力，后文会详细讲解。

DiT 通过 InputEmbedding 层（定义于 cosyvoice/flow/DiT/dit.py）将四路信号在入口处融合：

1
2
3
4
5
6
7
8
9
class InputEmbedding(nn.Module):
    def __init__(self, mel_dim=80, text_dim=80, out_dim=1024, spk_dim=80):
        self.proj = nn.Linear(80*2 + 80 + 80, 1024)  # = Linear(320, 1024)
        self.conv_pos_embed = CausalConvPositionEmbedding(dim=1024)

    def forward(self, x, cond, text_embed, spks):
        spks = repeat(spks, "b c -> b t c", t=x.shape[1])  # 广播到每一帧
        x = self.proj(torch.cat([x, cond, text_embed, spks], dim=-1))
        x = self.conv_pos_embed(x) + x  # 因果卷积位置编码 + 残差

拼接维度：80(噪声x) + 80(prompt mel) + 80(speech token embedding) + 80(说话人spk embedding) = 320 → Linear → 1024。

2.2 Prompt 训练/推理设计

CosyVoice 从 v1 到 v3 一直沿用相同的 prompt 条件设计，核心目的是让 Flow Matching 模型具备 in-context 能力：给定一段参考音频（prompt），生成与之音色、风格、音质一致的新语音。

2.2.1 训练阶段

训练时（cosyvoice/flow/flow.py 的 forward 方法），完整的 token 序列直接从 batch 中读入，不做任何 prompt / 非 prompt 的拆分或拼接。”prompt” 的概念仅通过条件变量 cond 体现，即随机选取序列前缀的一段 mel 帧作为条件输入：

1
2
3
4
5
6
conds = torch.zeros(feat.shape)
for i, j in enumerate(feat_len):
    if random.random() < 0.5:        # 50% 概率完全不给 prompt（cond 全零）
        continue
    index = random.randint(0, int(0.3 * j))  # 有 prompt 时，随机取前 0~30% 的 mel 帧
    conds[i, :index] = feat[i, :index]

DiT 的输入 y（即 noised mel）覆盖全序列，cond 作为额外条件信号与 y 一起 concat 输入 DiT。训练 loss 同样在全序列上计算（包含 prompt 区），prompt 区的梯度参与反向传播，使模型学会利用 prompt cond 中的参考信息来合成更好的效果。

针对 prompt 部分，cond 中包含 ground-truth mel 并不构成信息泄露：模型的预测目标是速度场 u = x₁ - (1 - σ_min) · z（其中x₁是真实梅尔），而非直接复原 mel 本身。即使 cond 提供了参考 mel，模型仍需从带噪状态 y 出发预测正确的速度方向。简而言之，cond 的作用是提供风格/音色参考，而非直接给出答案。

2.2.2 推理阶段

推理时（inference 方法），prompt 和待合成部分被显式分开传入，再拼接处理：

  • token 侧：prompt_token 和 token 通过 torch.concat 拼接成完整序列，经过 Embedding → PreLookahead → repeat 上采样，得到 token embedding mu
  • mel 条件侧：cond 的 prompt 区填入参考音频的真实 mel（prompt_feat），合成区为零——与训练时的 cond 结构完全对应
  • Flow Matching 过程：整条序列（含 prompt 区）从噪声出发，经 10 步 Euler ODE 推理
  • 输出裁剪：仅保留合成区的输出 mel（feat[:, :, mel_len1:]），丢弃 prompt 区的预测结果，送入后续声码器还原波形。

2.3 Token 编码与上采样

2.3.1 PreLookaheadLayer

PreLookaheadLayer（定义于 cosyvoice/transformer/upsample_encoder.py），是一个只有两层卷积的轻量 Encoder 模块，核心作用是为每个位置引入有限的未来信息（lookahead），同时保持整体的因果性。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
class PreLookaheadLayer(nn.Module):
    def __init__(self, in_channels=80, channels=1024, pre_lookahead_len=3):
        self.conv1 = nn.Conv1d(in_channels, channels, kernel_size=pre_lookahead_len + 1, ...)  # kernel=4, 非因果
        self.conv2 = nn.Conv1d(channels, in_channels, kernel_size=3, ...)                      # kernel=3, 因果

    def forward(self, inputs, context=torch.zeros(0, 0, 0)):
        outputs = inputs.transpose(1, 2).contiguous()           # [B, T, C] → [B, C, T]，转为 Conv1d 格式
        context = context.transpose(1, 2).contiguous()

        # ===== conv1: 非因果，需要右侧 lookahead =====
        if context.size(2) == 0:
            # 训练 / finalize 推理：右侧 pad 3 个零（无真实未来信息）
            outputs = F.pad(outputs, (0, self.pre_lookahead_len), mode='constant', value=0.0)
        else:
            # 流式推理非最后 chunk：拼接 3 个 context token 作为真实的未来信息
            # context.size(2) == pre_lookahead_len == 3，所以额外 pad 为 0
            outputs = F.pad(torch.concat([outputs, context], dim=2),
                            (0, self.pre_lookahead_len - context.size(2)), mode='constant', value=0.0)
        outputs = F.leaky_relu(self.conv1(outputs))
        # conv1 输出长度 = 输入长度（pad 后）- kernel_size + 1 = T + 3 - 4 + 1 = T，保持不变

        # ===== conv2: 因果，仅左侧 pad =====
        outputs = F.pad(outputs, (self.conv2.kernel_size[0] - 1, 0), mode='constant', value=0.0)  # 左侧 pad 2
        outputs = self.conv2(outputs)
        # conv2 输出长度 = T + 2 - 3 + 1 = T，保持不变

        outputs = outputs.transpose(1, 2).contiguous()          # [B, C, T] → [B, T, C]
        outputs = outputs + inputs                               # 残差连接
        return outputs

两层卷积的因果性设计并不一样：

  • conv1（kernel_size=4, padding=0）：输入右侧 pad 3 帧（零填充或传入 context），所以每个位置能看到自身 + 右侧 3 帧。这是整个模块中唯一的非因果操作，也是 “lookahead” 名称的由来
  • conv2（kernel_size=3, padding=0）：输入左侧 pad 2 帧（F.pad(outputs, (2, 0))），右侧 pad 0。这是标准的因果卷积，每个位置只看自身 + 左侧 2 帧，并不看未来信息

因此，只有 conv1 需要通过 context 机制处理未来信息，conv2 天然是因果的。推理时根据 finalize 区分两种模式：

场景 输入 context conv1 行为
非最后 chunk token[:, :-3] token[:, -3:] 拼接 context 后卷积，能看到后 3 个 token
最后 chunk（finalize=True） 全部 token 无，右侧 pad 零 与训练一致

这个设计保证了流式分 chunk 推理和全序列推理的输出逐帧完全一致，具体详见第 3 章关于流式的详细解读。

2.3.2 repeat_interleave 上采样

CosyVoice3 重新设计了帧率体系，使 token 帧率和 mel 帧率为精确的 2 倍整数关系：

  • token 帧率：25 Hz（token_frame_rate=25）
  • mel 帧率：24000 / 480 = 50 Hz（sample_rate=24000, hop_size=480）
  • 比值：token_mel_ratio = 2

每个 token 帧重复 2 次，能够直接匹配 mel 帧率，简化模型的结构和流式设计的难点。在严格两倍关系的基础上，上采样只需一行代码：

1
h = h.repeat_interleave(self.token_mel_ratio, dim=1)  # [B, T, 80] → [B, T×2, 80]

2.3.3 数据预处理保证帧数对齐

尽管 token 和梅尔特征的帧率在理论上是严格两倍的关系，但因为 token 和梅尔特征的提取细节存在细微差异，所以在部分特殊音频上，可能会触发两者相差1帧的问题（在语音生成任务上比较常见）。因此，为确保 token 数 × 2 == mel 帧数 严格成立，数据预处理阶段（cosyvoice/dataset/processor.py 中的 compute_fbank）会将音频长度规整到 960 的整数倍：

1
2
3
4
5
# cosyvoice/dataset/processor.py — compute_fbank
if num_frames != -1:  # num_frames=960
    index = int(np.ceil(sample['speech'].shape[1] / num_frames))
    sample['speech'] = torch.concat([sample['speech'],
        torch.zeros(1, index * num_frames - sample['speech'].shape[1])], dim=1)

960 = 24000/25 = 每个 token 对应的采样点数。pad 到 960 的倍数后：

  • mel 帧数 = L / 480 = 2k（偶数）
  • token 数 = L / 960 = k
  • token × 2 = mel 帧数，精确无误

2.4 DiT 网络

DiT（Diffusion Transformer，定义于 cosyvoice/flow/DiT/dit.py）是整个 flow matching 模块的核心。

2.4.1 模型配置参数

参数 值 含义
dim 1024 隐层维度
depth 22 Transformer 层数
heads 16 注意力头数
dim_head 64 每头维度（16×64=1024）
ff_mult 2 FFN 倍率（FFN 隐层=2048）
mel_dim 80 mel 特征维度
mu_dim 80 token embedding 维度
spk_dim 80 speaker embedding 投影后维度
out_channels 80 输出梅尔特征维度
static_chunk_size 50 流式 chunk 大小（梅尔特征的帧数，对应1秒）
num_decoding_left_chunks -1 默认使用全部左侧历史

2.4.2 模型结构与前向传播

1
2
3
4
5
6
7
8
9
10
11
class DiT(nn.Module):
    def __init__(self, *, dim, depth, heads, dim_head, ff_mult, mel_dim,
                 mu_dim, spk_dim, out_channels, static_chunk_size, num_decoding_left_chunks):
        self.time_embed = TimestepEmbedding(dim)                # 时间步嵌入
        self.input_embed = InputEmbedding(mel_dim, mu_dim, dim, spk_dim)  # 输入融合
        self.rotary_embed = RotaryEmbedding(dim_head)           # RoPE 位置编码
        self.transformer_blocks = nn.ModuleList(
            [DiTBlock(dim, heads, dim_head, ff_mult) for _ in range(depth)]
        )
        self.norm_out = AdaLayerNormZero_Final(dim)             # 最终自适应归一化
        self.proj_out = nn.Linear(dim, mel_dim)                 # 输出投影

前向传播流程：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
输入:
  x    [B, 80, T]  — 当前去噪状态（噪声 mel）
  mu   [B, 80, T]  — token 编码
  cond [B, 80, T]  — prompt mel 条件
  spks [B, 80]     — 说话人嵌入
  t    [B]         — 时间步

步骤:
  1. TimestepEmbedding:  t → SinusPos → MLP → [B, 1024]
  2. InputEmbedding:     cat(x, cond, mu, spks) [B, T, 320] → Linear → [B, T, 1024]
                         + CausalConvPositionEmbedding
  3. RoPE:               计算旋转位置编码
  4. 构造 attention mask（streaming/non-streaming，详见第 3 章）
  5. 22 × DiTBlock:      AdaLN-Zero(t调制) → Self-Attn(RoPE) → FFN
  6. AdaLN_Final(t调制) → Linear(1024 → 80)

输出:
  [B, 80, T]  — 预测的速度场 u = x₁ - (1 - σ_min) · z

2.4.3 DiTBlock 细节

每层 DiTBlock 的结构（定义于 cosyvoice/flow/DiT/modules.py），实际相当于在 Transformer Encoder 的基础上增加了时间 t 影响的 AdaLN-Zero。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
class DiTBlock(nn.Module):
    def __init__(self, dim, heads, dim_head, ff_mult, dropout):
        self.attn_norm = AdaLayerNormZero(dim)      # 自适应 LayerNorm
        self.attn = Attention(AttnProcessor(), dim, heads, dim_head, dropout)
        self.ff_norm = nn.LayerNorm(dim, elementwise_affine=False)
        self.ff = FeedForward(dim, mult=ff_mult, approximate="tanh")

    def forward(self, x, t, mask=None, rope=None):
        # 1. AdaLN-Zero：时间步 t 生成 6 个调制参数
        norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t)
        # 2. Self-Attention with RoPE
        attn_output = self.attn(x=norm, mask=mask, rope=rope)
        x = x + gate_msa.unsqueeze(1) * attn_output
        # 3. FFN with adaptive modulation
        ff_norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
        x = x + gate_mlp.unsqueeze(1) * self.ff(ff_norm)

其中，AdaLN-Zero（cosyvoice/flow/DiT/modules.py 中的 AdaLayerNormZero）是 DiT 将时间步 t 注入模型的核心机制，源自 DiT 原始论文。与标准 LayerNorm 使用固定的可学习 affine 参数不同，AdaLN-Zero 的归一化参数由时间步 t 关联生成（这种关联/影响，可以称为「调制」，类似 AdaSpeech 系列的 Conditional LayerNorm 设计思想）：

1
2
3
4
5
6
7
8
9
10
11
class AdaLayerNormZero(nn.Module):
    def __init__(self, dim):
        self.silu = nn.SiLU()
        self.linear = nn.Linear(dim, dim * 6)  # 一次性生成 6 个调制向量
        self.norm = nn.LayerNorm(dim, elementwise_affine=False)  # 无可学习 affine 参数

    def forward(self, x, emb=None):
        emb = self.linear(self.silu(emb))  # emb 即时间步嵌入 t, shape: [B, dim]
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(emb, 6, dim=1)
        x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]  # 自适应归一化
        return x, gate_msa, shift_mlp, scale_mlp, gate_mlp

6 个调制向量均来自于 t 投影后的 embedding，按照作用的模块分为两组，分别用于 Attention 和 FFN 子层：

调制参数 作用位置 用途
scale_msa / shift_msa Attention 前 对 LayerNorm 输出做仿射变换：LN(x) * (1 + scale) + shift
gate_msa Attention 后 门控 Attention 输出：x = x + gate_msa * attn_output
scale_mlp / shift_mlp FFN 前 对 FFN 输入的 LayerNorm 做仿射变换
gate_mlp FFN 后 门控 FFN 输出：x = x + gate_mlp * ff_output

设计要点：

  • self.norm 设置 elementwise_affine=False，是为了移除标准 LayerNorm 的可学习 γ/β 参数，改为完全由时间步 t embedding 控制，相当于模型可以学习到不同时间t的区别
  • AdaLN-Zero 名称中的 “Zero”，原本意思是将self.linear = nn.Linear(dim, dim * 6)线性层的权重初始化为零，使训练初期 DiTBlock 近似恒等映射，有利于深层网络训练稳定性（但在当前 CosyVoice 代码实现中，没有修改，nn.Linear 沿用了 PyTorch 默认的 Kaiming Uniform 初始化）
  • 最后一层 DiTBlock 之后使用的是 AdaLayerNormZero_Final，仅生成 scale、shift，因为后面没有 Attention 和 FFN 层了。

2.4.4 因果卷积位置编码

DiT 使用因果卷积（CausalConvPositionEmbedding）代替传统的双向卷积位置编码（均定义于 cosyvoice/flow/DiT/modules.py）：

1
2
3
4
5
6
7
8
9
10
class CausalConvPositionEmbedding(nn.Module):
    def __init__(self, dim, kernel_size=31, groups=16):
        self.conv1 = nn.Sequential(nn.Conv1d(dim, dim, 31, groups=16, padding=0), nn.Mish())
        self.conv2 = nn.Sequential(nn.Conv1d(dim, dim, 31, groups=16, padding=0), nn.Mish())

    def forward(self, x):
        x = F.pad(x, (30, 0))  # 仅左侧 padding
        x = self.conv1(x)
        x = F.pad(x, (30, 0))  # 仅左侧 padding
        x = self.conv2(x)

与非因果版 ConvPositionEmbedding（两侧 padding=15）不同，因果版只在左侧 pad 30 帧，右侧 pad 0，是严格因果的结构。每个位置仅依赖当前及之前 30 帧的信息，这是支持流式推理的基础设计。

2.5 训练目标

Flow Matching（代码位于 cosyvoice/flow/flow_matching.py）定义了一个从噪声分布 到数据分布 的概率路径。给定目标样本 （真实 mel）和噪声 ，在时间 上的插值为：

对应的目标速度场为：

模型（DiT）正是基于传入的各种条件信息，被训练来预测这个速度场，损失为 MSE：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# CausalConditionalCFM.compute_loss
t = torch.rand([b, 1, 1])                           # 随机时间步 ∈ [0, 1]
z = torch.randn_like(x1)                            # 随机噪声
y = (1 - (1 - sigma_min) * t) * z + t * x1          # 插值
u = x1 - (1 - sigma_min) * z                        # 目标速度场

# CFG 训练：20% 概率 drop 全部条件（training_cfg_rate=0.2）
if training_cfg_rate > 0:
    cfg_mask = torch.rand(b) > 0.2
    mu   = mu   * cfg_mask.view(-1, 1, 1)
    spks = spks * cfg_mask.view(-1, 1)
    cond = cond * cfg_mask.view(-1, 1, 1)

pred = self.estimator(y, mask, mu, t, spks, cond)   # DiT 预测速度场
loss = MSE(pred * mask, u * mask)                     # 全序列 loss

训练过程中，training_cfg_rate 设置 20% 概率将 mu / spks / cond 全部置零，让模型同时学习有条件和无条件分布，是 Classifier-Free Guidance（CFG）的训练配置。

2.6 推理采样

推理时，从固定噪声出发，经过 10 步 Euler ODE 求解。完整流程如下（代码位于 cosyvoice/flow/flow_matching.py）：

2.6.1 Cosine 时间步调度

1
2
3
4
5
# CausalConditionalCFM.forward
z = self.rand_noise[:, :, :mu.size(2)].to(mu.device).to(mu.dtype) * temperature
t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype)
if self.t_scheduler == 'cosine':
    t_span = 1 - torch.cos(t_span * 0.5 * torch.pi)

Cosine 调度后的时间步分布：

步骤 t 区间 dt 特点
1 0.000 → 0.025 0.025 高噪声区，小步精细去噪
2-4 0.025 → 0.345 递增 逐步加速
5-6 0.345 → 0.655 ~0.155 中间阶段步长最大
7-9 0.655 → 0.975 递减 逐步减速
10 0.975 → 1.000 0.025 低噪声区，小步精细打磨

2.6.2 CFG 推理与 Euler 求解

每步 Euler 迭代中，构造 batch=2 同时做有条件和无条件推理，代码简写为：

1
2
3
4
5
6
7
8
9
10
11
12
# 10 步 Euler 循环
for step in range(10):
    # batch[0]: 有条件 — mu/spks/cond 有值
    # batch[1]: 无条件 — mu/spks/cond 全零
    dphi_dt = DiT(x_in, ...)  # [2, 80, T]
    dphi_cond, dphi_uncond = split(dphi_dt)

    # CFG 公式: inference_cfg_rate=0.7
    velocity = (1 + 0.7) * dphi_cond - 0.7 * dphi_uncond
             = 1.7 * dphi_cond - 0.7 * dphi_uncond

    x = x + dt * velocity    # Euler 更新

cfg_rate=0.7 意味着对有条件预测做 1.7 倍放大、对无条件预测做 0.7 倍抑制，增强条件引导的效果。每步需要执行一次 DiT forward（batch=2 用于 CFG），10 步共 10 次 forward、20 次 DiT 计算。

以上就是 DiT 模型从输入信息到模型结构、从训练目标和推理方式的解读，结合 DiT 和 Flow Matching 相关的论文能够更方便理解。

3. 流式推理设计

CosyVoice3 DiT 最重要的参考细节是在流式的设计上，在同一个模型、同一套参数下，通过不同的 attention mask 同时支持流式和非流式推理。具体包括以下四点细节：

1
2
3
4
细节 1: Unified Training    50% 概率 streaming / 非 streaming 联合训练
细节 2: PreLookahead        因果卷积 + context 前瞻机制
细节 3: DiT Attention       chunk mask vs full mask
细节 4: CFM 固定噪声         预生成噪声保证 chunk 间一致

3.1 Unified Training

1
2
# CausalMaskedDiffWithDiT.forward (训练)
streaming = True if random.random() < 0.5 else False

训练时每个 batch 有 50% 概率使用 chunk mask、50% 使用 full mask。streaming 实际最终影响的是 DiT 的 attention mask 形状：

  • streaming=True：chunk mask（受限视野）
  • streaming=False：full attention（全局视野）

3.2 Chunk Attention Mask

3.2.1 subsequent_chunk_mask

该函数定义于 cosyvoice/utils/mask.py，负责生成 chunk 级的 causal mask：

1
2
3
4
5
def subsequent_chunk_mask(size, chunk_size, num_left_chunks=-1, device='cpu'):
    pos_idx = torch.arange(size, device=device)
    block_value = (torch.div(pos_idx, chunk_size, rounding_mode='trunc') + 1) * chunk_size
    ret = pos_idx.unsqueeze(0) < block_value.unsqueeze(1)
    return ret

以 size=8, chunk_size=2 为例：

1
2
3
4
5
6
7
8
9
位置    0  1  2  3  4  5  6  7
  0  [  1  1  0  0  0  0  0  0 ]   ← chunk 0 只看 [0,1]
  1  [  1  1  0  0  0  0  0  0 ]
  2  [  1  1  1  1  0  0  0  0 ]   ← chunk 1 看 [0..3]
  3  [  1  1  1  1  0  0  0  0 ]
  4  [  1  1  1  1  1  1  0  0 ]   ← chunk 2 看 [0..5]
  5  [  1  1  1  1  1  1  0  0 ]
  6  [  1  1  1  1  1  1  1  1 ]   ← chunk 3 看全部
  7  [  1  1  1  1  1  1  1  1 ]

可以直观地看出，chunk mask 保证的是同一 chunk 内的帧互相可见，且能看到之前所有 chunk。

3.2.2 chunk_mask 实现细节

DiT 代码中（cosyvoice/flow/DiT/dit.py），根据 streaming 标志选择不同的 mask 策略：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
# DiT.forward
if streaming is True:
    attn_mask = add_optional_chunk_mask(
        x, mask.bool(),
        use_dynamic_chunk=False, use_dynamic_left_chunk=False,
        decoding_chunk_size=0,
        static_chunk_size=self.static_chunk_size,  # 50
        num_decoding_left_chunks=-1      # 看全部历史
    ).unsqueeze(dim=1)       # (B, L, L) → (B, 1, L, L) 供多头 attention 广播
else:
    attn_mask = add_optional_chunk_mask(
        x, mask.bool(),
        use_dynamic_chunk=False, use_dynamic_left_chunk=False,
        decoding_chunk_size=0,
        static_chunk_size=0,  # 0 → 不做 chunk 分块
        num_decoding_left_chunks=-1
    ).repeat(1, x.size(1), 1).unsqueeze(dim=1)     # (B, 1, L) → (B, L, L) → (B, 1, L, L)

关键区别在于 static_chunk_size 和后续的 shape 处理：

Streaming（static_chunk_size=50）：进入 add_optional_chunk_mask 的 elif static_chunk_size > 0 分支，调用 subsequent_chunk_mask(L, 50, -1) 生成 (L, L) 的 chunk 因果 mask，每个位置只能看到自身所在 chunk 及之前所有 chunk 的帧。

Non-streaming（static_chunk_size=0）：进入 else 分支，直接返回 padding mask (B, 1, L)。然后通过 .repeat(1, L, 1) 扩展为 (B, L, L)，每一行都是相同的 padding mask，意味着所有非 padding 位置之间双向可见，相当于全局 attention。

3.3 PreLookahead 的 context 机制

之前在 2.3 小节中提到，PreLookaheadLayer 需要看到右侧 3 帧。流式推理中，每次调用时额外传入后 3 个 token 作为 context（代码位于 cosyvoice/flow/flow.py）：

1
2
3
4
if finalize is True: # 最后 chunk: 无 context，pad 零（与训练一致）
    h = self.pre_lookahead_layer(token)
else: # 非最后 chunk: 多传 3 个 token 作 context
    h = self.pre_lookahead_layer(token[:, :-self.pre_lookahead_len], context=token[:, -self.pre_lookahead_len:])

context 虽然参与了 conv1 的卷积计算，但不产生输出，这样能保证每个位置的卷积范围内容，和完整序列处理时完全一致。

3.4 固定噪声

CausalConditionalCFM（定义于 cosyvoice/flow/flow_matching.py）使用预生成的固定噪声，代替完全的随机采样：

1
2
3
4
5
6
7
class CausalConditionalCFM(ConditionalCFM):
    def __init__(self, ...):
        set_all_random_seed(0)
        self.rand_noise = torch.randn([1, 80, 50 * 300])  # 预生成，覆盖 300 秒

    def forward(self, mu, ...):
        z = self.rand_noise[:, :, :mu.size(2)]  # 截取所需长度

固定噪声解决的问题：流式推理中会多次调用 forward，如果每次随机采样，同一个位置在不同次调用中的初始噪声不一致，是不符合 Flow Matching 的原理的。而使用固定噪声后，位置 i 的噪声永远是 rand_noise[:,:,i]。

注意：训练时使用随机噪声 torch.randn_like(x1)，不需要固定，因为训练是一次处理完整序列，流式是通过 attn_mask 的设计来实现的，并不存在多次调用的一致性问题。

3.5 保证推理一致性

综合上面的信息，以下多个机制共同保证分 chunk 推理和全序列推理结果严格相等（精度误差为 0）：

  1. PreLookahead context → 卷积输出逐位置一致
  2. repeat_interleave → 这一步是确定性的上采样
  3. 固定噪声 → 每个位置的初始噪声值在不同次推理中相同
  4. Chunk attention mask → 每个位置的注意力范围在不同次推理中完全相同（num_left_chunks=-1，看全部历史）
  5. 因果卷积位置编码 → 不依赖未来，输出也逐位置一致

代码中附带了验证脚本（cosyvoice/flow/flow.py 末尾的 __main__ 块），分 chunk 推理后逐帧对比全序列推理结果，差异为 0。

3.6 流式推理的完整实现

上面 3.1~3.5 讲的是 DiT 模型层面的流式设计，本节从工程实践角度展示完整的流式推理链路：LLM 异步产 token → Flow Matching 分 chunk 生成 mel → HiFi-GAN 合成波形。代码位于 cosyvoice/cli/model.py 的 CosyVoice3Model（继承自 CosyVoice2Model）。

3.6.1 token2wav 内部细节

CosyVoice3Model 的初始化和核心方法 token2wav：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
class CosyVoice3Model(CosyVoice2Model):

    def __init__(self, llm, flow, hift, fp16=False):
        # ...
        # token_hop_len: 每轮 Flow 推理，需要的新增输出的 token 数（不含 context）
        # 注意：DiT 每次推理仍输入全部历史 token，但只截取最后 token_hop_len 个 token 对应的新增 mel
        self.token_hop_len = 25                             # 初始值 = 训练的 chunk_size，对应 50 mel 帧 = 1 秒
        # token_max_hop_len: hop 的上限，平衡调用次数与输出间隔（hop 太大会导致两次音频推理，等待输出的时间过久）
        self.token_max_hop_len = 4 * self.token_hop_len     # = 100 token，对应 200 mel 帧 = 4 秒
        # stream_scale_factor: hop_len 的增长倍率，每轮翻倍
        # 首包 25 tok (低延迟) → 第2包 50 tok → 第3包起每新增 100 tok 推理一次
        self.stream_scale_factor = 2

    def token2wav(self, token, prompt_token, prompt_feat, embedding,
                  token_offset, uuid, stream=False, finalize=False, speed=1.0):
        with torch.cuda.amp.autocast(self.fp16):
            # ===== 1. Flow Matching 推理：token → mel =====
            # token 包含 prompt_token 之后的全部已产出 token（含 context）
            # flow.inference 内部会拼接 prompt_token，去噪后丢弃 prompt 区 mel
            tts_mel, _ = self.flow.inference(
                token=token.to(self.device, dtype=torch.int32),
                token_len=torch.tensor([token.shape[1]], dtype=torch.int32).to(self.device),
                prompt_token=prompt_token.to(self.device),
                prompt_token_len=torch.tensor([prompt_token.shape[1]], dtype=torch.int32).to(self.device),
                prompt_feat=prompt_feat.to(self.device),
                prompt_feat_len=torch.tensor([prompt_feat.shape[1]], dtype=torch.int32).to(self.device),
                embedding=embedding.to(self.device),
                streaming=stream, finalize=finalize)
            # 返回的 tts_mel 是全部合成区 mel（不含 prompt），shape: [1, 80, total_mel_len]

            # ===== 2. 裁剪 mel：只保留当前 chunk 新增的帧 =====
            # token_offset 是之前已输出的 token 数，乘以 token_mel_ratio=2 换算为 mel 帧偏移
            tts_mel = tts_mel[:, :, token_offset * self.flow.token_mel_ratio:]

            # ===== 3. 拼接 HiFi-GAN 缓存 mel =====
            # HiFi-GAN 需要一定的历史 mel 才能保证 chunk 衔接处波形连续
            if self.hift_cache_dict[uuid] is not None:
                hift_cache_mel = self.hift_cache_dict[uuid]['mel']
                tts_mel = torch.concat([hift_cache_mel, tts_mel], dim=2)
                self.hift_cache_dict[uuid]['mel'] = tts_mel          # 更新缓存为累积 mel
            else:
                self.hift_cache_dict[uuid] = {'mel': tts_mel, 'speech_offset': 0}

            # ===== 4. HiFi-GAN 声码器：mel → waveform =====
            tts_speech, _ = self.hift.inference(speech_feat=tts_mel, finalize=finalize)
            # 通过 speech_offset 截取新增波形，避免重复输出
            tts_speech = tts_speech[:, self.hift_cache_dict[uuid]['speech_offset']:]
            self.hift_cache_dict[uuid]['speech_offset'] += tts_speech.shape[1]
        return tts_speech

token2wav 的核心逻辑：

  • flow.inference 内部已丢弃 prompt 区 mel（feat[:, :, mel_len1:]），返回全部合成区 mel
  • token_offset * token_mel_ratio 进一步裁掉之前 chunk 已输出的部分，只留当前 chunk 新增的 mel 帧
  • HiFi-GAN 通过 speech_offset 记录已输出的波形长度，每次只截取新增波形，这样就不需要 CosyVoice2 中的 overlap + fade-in/out 平滑拼接，因为 DiT 流式设计已保证 chunk 边界处的 mel 严格连续

3.6.2 流式 TTS 调用流程

tts 方法（继承自 CosyVoice2Model）协调 LLM 和 Flow Matching 的流水线：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
# CosyVoice2Model.tts —— 流式分支

# ===== 启动 LLM 线程，异步产出 speech token =====
p = threading.Thread(target=self.llm_job, args=(..., this_uuid))
p.start()   # LLM 持续将 token append 到 self.tts_speech_token_dict[this_uuid]

token_offset = 0
# 首 chunk 需要将 prompt token 长度对齐到 token_hop_len 的整数倍
# 确保首 chunk 的 token 数（含 prompt pad）恰好让 mel 帧对齐 static_chunk_size=50 的边界
prompt_token_pad = int(np.ceil(flow_prompt_speech_token.shape[1] / self.token_hop_len)
                       * self.token_hop_len - flow_prompt_speech_token.shape[1])

while True:
    time.sleep(0.1)   # 轮询间隔 100ms
    # 首 chunk 的 hop 需要加上 prompt 对齐 pad，后续 chunk 直接用 token_hop_len
    this_token_hop_len = self.token_hop_len + prompt_token_pad if token_offset == 0 else self.token_hop_len

    # ===== 判断 LLM 已产出的 token 是否够一次 Flow 推理 =====
    # 需要 hop_len 个输出 token + pre_lookahead_len(3) 个 context token
    if len(self.tts_speech_token_dict[this_uuid]) - token_offset >= this_token_hop_len + self.flow.pre_lookahead_len:
        # 取从头到 (token_offset + hop + 3) 的全部 token，含 context
        this_tts_speech_token = torch.tensor(
            self.tts_speech_token_dict[this_uuid][:token_offset + this_token_hop_len + self.flow.pre_lookahead_len]
        ).unsqueeze(dim=0)
        this_tts_speech = self.token2wav(token=this_tts_speech_token,
                                         prompt_token=flow_prompt_speech_token,
                                         prompt_feat=prompt_speech_feat,
                                         embedding=flow_embedding,
                                         token_offset=token_offset,
                                         uuid=this_uuid, stream=stream, finalize=False)
        token_offset += this_token_hop_len
        # 渐进增大 hop: 25 → 50 → 100 (上限)，首包低延迟，后续重视吞吐
        # 核心约束：RTF < 1，即每轮生成耗时 < 该 chunk 对应的音频时长
        # 首包 1 秒音频播放期间，系统有"喘息时间"去准备更大的后续 chunk，如此滚动不造成流式播放的卡顿
        self.token_hop_len = min(self.token_max_hop_len, self.token_hop_len * self.stream_scale_factor)
        yield {'tts_speech': this_tts_speech.cpu()}

    # LLM 已结束 且 剩余 token 不够凑一个完整 chunk → 跳出循环进入 finalize
    if self.llm_end_dict[this_uuid] is True and \
       len(self.tts_speech_token_dict[this_uuid]) - token_offset < this_token_hop_len + self.flow.pre_lookahead_len:
        break

p.join()
# ===== 最后一次 finalize 推理：处理全部剩余 token，无 context =====
this_tts_speech_token = torch.tensor(self.tts_speech_token_dict[this_uuid]).unsqueeze(dim=0)
this_tts_speech = self.token2wav(token=this_tts_speech_token, ..., token_offset=token_offset, finalize=True)
yield {'tts_speech': this_tts_speech.cpu()}

从代码中可以看到几个关键设计：

  1. 等待条件 >= hop_len + pre_lookahead_len：每次送入 Flow 的 token 比实际输出多 3 个，这 3 个就是 PreLookahead 的 context（参见 3.3 节），参与卷积但不产生输出
  2. 首 chunk 对齐：prompt_token_pad 将 prompt token 长度向上取整到 token_hop_len=25 的倍数，保证 chunk 边界与 DiT 的 static_chunk_size=50 mel 帧对齐
  3. 渐进式增大 hop：token_hop_len 每次乘以 stream_scale_factor=2，从 25→50→100（上限），首包用最小 chunk 降低延迟，后续用更大 chunk 提高推理吞吐
  4. finalize 收尾：LLM 结束后，剩余 token 不足一个 hop 时跳出循环，最后一次调用 token2wav(finalize=True) 处理全部剩余 token

以 3 个 chunk 为例（chunk_size=25 token, pre_lookahead_len=3, token_mel_ratio=2），完整流程如图：

3.8 关于 Dynamic Chunk 的补充

cosyvoice/utils/mask.py 中的 add_optional_chunk_mask 函数还支持两层 dynamic 随机化（源自 WeNet 流式 ASR 框架），通过 use_dynamic_chunk 和 use_dynamic_left_chunk 两个开关分别控制。CosyVoice3 的 flow 模块未开启此模式，但是可以自己进行实验的。

3.8.1 Dynamic Chunk（动态 chunk 大小）

设置 use_dynamic_chunk=True 开启。训练时（decoding_chunk_size=0），每个 batch 随机采样 chunk 大小：

1
2
3
4
5
chunk_size = torch.randint(1, max_len, (1,)).item()
if chunk_size > max_len // 2 and enable_full_context:
    chunk_size = max_len        # ~50% 概率：全序列（等价于非流式）
else:
    chunk_size = chunk_size % 25 + 1   # ~50% 概率：chunk ∈ [1, 25]

核心思想：模型在训练时见过各种 chunk 大小，推理时可以指定任意 decoding_chunk_size，灵活适配不同延迟需求，chunk 越大延迟越高但质量越好，chunk 越小延迟越低。

3.8.2 Dynamic Left Chunk（动态左侧历史 chunk 个数）

在 Dynamic Chunk 基础上，进一步设置 use_dynamic_left_chunk=True 开启。控制的是每个位置在 attention 中能回看多少个历史 chunk：

1
2
3
if use_dynamic_left_chunk:
    max_left_chunks = (max_len - 1) // chunk_size
    num_left_chunks = torch.randint(0, max_left_chunks, (1,)).item()

不开启时 num_left_chunks=-1，即看全部历史。开启后，训练时随机限制可回看的历史 chunk 数量（从 0 到最大值均匀采样），推理时可以通过 num_decoding_left_chunks 指定固定值。这解决的是显存/计算量随序列增长线性膨胀的问题，限制左侧窗口后，适合超长音频的流式场景。为了保证模型效果的稳定性，训练时开启 use_dynamic_left_chunk 对于推理时限制长度理论上也是有好处的。

use_dynamic_chunk use_dynamic_left_chunk chunk 大小 左侧历史 适用场景
False False 固定 (static_chunk_size) 全部 (-1) CosyVoice3 当前方案，chunk 固定，左侧窗长持续增长
True False 随机 [1,25] 或全序列 全部 (-1) 灵活延迟，但左侧窗长持续增长
True True 随机 [1,25] 或全序列 随机 [0, max] 灵活延迟 + 控制显存

CosyVoice3 开源的是最简洁的方案：static chunk（固定 50 mel 帧），不做 chunk 大小或左侧窗口的动态随机化，随机性只出现在 50% 概率 streaming/non-streaming 随机切换。

4. 简要总结

CosyVoice3 的 Flow Matching + DiT 模块代表了 TTS 设计的几个趋势：

  1. Transformer 取代 U-Net：22 层 DiT 替代了传统 U-Net decoder，更适合序列建模且易于流式化
  2. 极简编码器：两层因果卷积（PreLookaheadLayer）+ repeat 上采样，模型参数更集中在 DiT 模块。
  3. 训练-推理一致的流式设计：通过 PreLookahead context、固定噪声、chunk mask 三者配合，实现了可靠的流式/非流式推理
  4. 统一训练：50% streaming / 50% non-streaming 随机切换，一套权重适配两种场景

Flow Matching / DiT 的基础原理、模型设计、流式处理，还是存在很多技术细节的，是一个生成式算法和工程设计精妙结合的典范。虽然近来 TTS 的一部分新方案在逐步去除 DiT（比如重新回到 LLM + RVQ 的思路），但其最终是否能够真正被替代还需要实践验证。

  • 本文标题：代码解读 | CosyVoice 代码研读（三）：CosyVoice DiT CFM
  • 本文作者：白御青
  • 创建时间：2026-01-29
  • 本文链接：2026/2026-01-29-cosyvoice3-dit/
  • 版权声明：本博客所有文章除特别声明外，均采用 BY-NC-SA 许可协议。转载请注明出处！
  1. 1. 整体架构概览
  2. 2. 模型结构解读
    1. 2.1 输入条件与信息融合
    2. 2.2 Prompt 训练/推理设计
      1. 2.2.1 训练阶段
      2. 2.2.2 推理阶段
    3. 2.3 Token 编码与上采样
      1. 2.3.1 PreLookaheadLayer
      2. 2.3.2 repeat_interleave 上采样
      3. 2.3.3 数据预处理保证帧数对齐
    4. 2.4 DiT 网络
      1. 2.4.1 模型配置参数
      2. 2.4.2 模型结构与前向传播
      3. 2.4.3 DiTBlock 细节
      4. 2.4.4 因果卷积位置编码
    5. 2.5 训练目标
    6. 2.6 推理采样
      1. 2.6.1 Cosine 时间步调度
      2. 2.6.2 CFG 推理与 Euler 求解
  3. 3. 流式推理设计
    1. 3.1 Unified Training
    2. 3.2 Chunk Attention Mask
      1. 3.2.1 subsequent_chunk_mask
      2. 3.2.2 chunk_mask 实现细节
    3. 3.3 PreLookahead 的 context 机制
    4. 3.4 固定噪声
    5. 3.5 保证推理一致性
    6. 3.6 流式推理的完整实现
      1. 3.6.1 token2wav 内部细节
      2. 3.6.2 流式 TTS 调用流程
    7. 3.8 关于 Dynamic Chunk 的补充
      1. 3.8.1 Dynamic Chunk（动态 chunk 大小）
      2. 3.8.2 Dynamic Left Chunk（动态左侧历史 chunk 个数）
  4. 4. 简要总结