Jekyll2026-06-11T01:29:19+00:00https://wqh011128.github.io/feed.xmlQihang WuTraining and inference framework engineer focused on LLM and VLM systems, with earlier research on domain adaptation for semantic segmentation and efficient attention mechanisms.Pallas TPU Kernel 写法与优化文档2026-06-10T00:00:00+00:002026-06-10T00:00:00+00:00https://wqh011128.github.io/blog/2026/06/10/Pallas_Kernel%E5%AD%A6%E4%B9%A0%E6%96%87%E6%A1%A3Pallas TPU Kernel 写法与优化文档

0. 适用范围

本文总结 Pallas TPU kernel 的常用生产级写法、设计规范与优化思路，重点面向以下类型的 kernel：

  • Attention / FlashAttention / GQA / MQA；
  • Matmul-heavy backward kernel；
  • 带 reduction 的 block kernel；
  • 需要在 TPU MXU 上获得较高利用率的 Pallas kernel；
  • 需要控制 HBM、VMEM、register、scratch、tile shape 的性能敏感 kernel。

本文尤其关注：

1. 如何设计 grid
2. 如何设计 BlockSpec
3. 如何使用 scratch 做 partial accumulation
4. 如何选择 major/minor tile
5. 如何避免单个 program 内部 loop 太长
6. 如何提升 MXU utilization
7. 如何 benchmark 和调 block size
8. 如何写出更接近生产级的 Pallas kernel

1. Pallas Kernel 的基本心智模型

写 Pallas kernel 时，建议始终从下面几个概念出发。

1.1 Program

一个 Pallas program 可以理解成一个小的 tile 程序实例。

它由 grid 决定数量，例如：

grid = (
    batch_blocks,
    num_heads,
    q_blocks,
    kv_blocks,
)

每个 program 会通过：

pl.program_id(axis)

拿到当前自己负责的 tile index。

例如：

q_block_index = pl.program_id(2)
kv_block_index = pl.program_id(3)

可以理解为：

一个 program 负责一个局部 tile。

1.2 Grid

grid 决定有多少个 program 被调度。

一个好的 grid 应该：

1. 暴露足够并行度；
2. 把重要的 block 维度显式暴露给 compiler；
3. 避免把过长的 reduction loop 藏在单个 program 里；
4. 让每个 program 的工作量适中；
5. 让输出 tile 和 reduction tile 的关系清楚。

对于 attention backward，常见 grid 设计是：

DQ:
  grid = batch × head × q_block × kv_block

DKV:
  grid = batch × head_or_kv_head × kv_block × q_block

1.3 BlockSpec

pl.BlockSpec 决定每个 program 看到输入/输出 tensor 的哪一个 tile。

典型形式：

q_spec = pl.BlockSpec(
    (block_b, group_size, block_q, head_dim),
    q_index_map,
)

其中：

block_shape:
  当前 program 看到的局部 tile shape

index_map:
  从 program_id 映射到全局 tensor block index

生产级原则：

BlockSpec 应该尽量描述局部 tile，而不是 full sequence。

例如，尽量避免：

q_spec = pl.BlockSpec(
    (block_b, group_size, q_seq_len, head_dim),
    q_full_index_map,
)

因为这意味着一个 program 对 Q 的 sequence 维几乎没有 tile 化。

更推荐：

q_spec = pl.BlockSpec(
    (block_b, group_size, block_q_major, head_dim),
    q_index_map,
)

1.4 MXU

TPU MXU 是主要做矩阵乘法的硬件单元。

Pallas kernel 想跑得快，通常要做到：

1. dot_general / dot 的 shape 规整；
2. matmul tile 足够大；
3. matmul 之间不要被太多 elementwise / mask / exp 阻塞；
4. 数据能及时 load 到 VMEM；
5. 不要让 MXU 等待数据或等待复杂控制流；
6. 不要让单个 program 的 live range 太长。

如果 MXU utilization 很低，例如只有 13%，通常说明：

MXU 大量时间没有被有效喂饱。

常见原因包括：

1. grid 并行度不足；
2. 单个 program 内部 loop 太长；
3. temporary tensor 太大；
4. VMEM/register 压力过高；
5. BlockSpec 过大；
6. dot tile shape 不合适；
7. matmul 中间夹杂过多 exp/mask/elementwise；
8. causal 上三角无效 block 没有跳过。

2. 写 Kernel 前必须先做的设计

不要一上来就写 Pallas kernel。生产级 kernel 应先写清楚：

1. 数学公式；
2. 输入输出 shape；
3. 哪些维度是并行维；
4. 哪些维度是 reduction 维；
5. 是否需要跨 block 累加；
6. 是否需要 scratch；
7. 是否需要保存 residuals；
8. 是否需要 mask/bias/causal/segment；
9. 是否需要支持 padding；
10. 是否需要支持 GQA/MQA。

2.1 先写数学公式

以 attention backward 为例：

前向：

S = QK^T
P = softmax(S)
O = P V

反向：

dV = P^T dO
dP = dO V^T
D  = sum(O * dO, axis=-1)
dS = P * (dP - D)
dQ = dS K
dK = dS^T Q

如果有 scale：

S = QK^T * sm_scale

则：

dS_unscaled = dS_scaled * sm_scale

如果有 additive attention bias：

S = QK^T + AB

则：

dAB = dS

2.2 明确输出决定 kernel 方向

对于 backward，常见拆法是：

DQ kernel:
  输出 dQ
  沿 KV 维 reduction

DKV kernel:
  输出 dK, dV
  沿 Q 维 reduction

因此：

DQ 的 grid 应该优先围绕 Q output tile 设计；
DKV 的 grid 应该优先围绕 KV output tile 设计。

3. Grid 设计规范

3.1 不要把长 reduction 全塞进单个 program

不推荐：

grid = (
    batch_blocks,
    num_heads,
    q_blocks,
)

for kv_start in range(0, kv_seq_len, block_kv):
    ...

因为这意味着：

一个 program 负责一个 q_block，
然后在 program 内部扫完整 kv_seq。

如果 kv_seq_len 很长，program 内部 loop 会很长。

更推荐：

grid = (
    batch_blocks,
    num_heads,
    q_blocks,
    kv_blocks,
)

让 kv_blocks 成为 grid 维度。

3.2 DQ 推荐 grid

DQ 的公式：

dQ = sum_over_KV dS @ K

推荐：

grid_dq = (
    batch_blocks,
    num_kv_heads,
    q_seq_len // block_q_major_dq,
    kv_seq_len // block_kv_major_dq,
)

含义：

axis 0: batch block
axis 1: kv head
axis 2: q major block
axis 3: kv major block

如果是 GQA：

一个 kv_head 对应 group_size 个 q_heads。

所以 Q tile 可以是：

[block_b, group_size, block_q_major_dq, head_dim]

3.3 DKV 推荐 grid

DKV 的公式：

dK = sum_over_Q dS^T @ Q
dV = sum_over_Q P^T @ dO

推荐：

grid_dkv = (
    batch_blocks,
    num_kv_heads,
    kv_seq_len // block_kv_major_dkv,
    q_seq_len // block_q_major_dkv,
)

含义：

axis 0: batch block
axis 1: kv head
axis 2: kv major block
axis 3: q major block

这样每个 program 负责：

一个 KV output tile
一个 Q reduction tile

再通过 scratch 跨 Q block 累加到最终 dK/dV。

3.4 哪些维度放 grid，哪些维度放 loop？

经验规则：

短 loop 可以放 program 内部；
长 reduction loop 尽量暴露到 grid；
需要跨 block 累加时，用 scratch；
如果 scratch/ordering 很难处理，再考虑内部 loop。

推荐：

batch/head/q_block/kv_block:
  通常适合放 grid

minor tile:
  通常适合放 program 内部 loop

完整 seq_len:
  不建议放单个 program 内部扫完

4. Scratch 设计规范

4.1 Scratch 的作用

Scratch 不是为了“直接加速”，而是为了允许：

把 reduction 维拆到 grid
+
在 VMEM 中保存 partial accumulation
+
避免 atomic
+
最后一个 reduction block 写出结果

如果没有 scratch，通常只能让一个 program 扫完整 reduction 维。

这容易造成：

1. 单个 program 太重；
2. 内部 loop 太长；
3. compiler 调度困难；
4. live range 变长；
5. MXU 利用率低。

4.2 DQ scratch 模式

DQ 需要沿 KV 维累加：

dQ = sum_over_KV dS @ K

所以 DQ scratch 可以设计为：

dq_scratch = pltpu.VMEM(
    (block_b, group_size, block_q_major_dq, head_dim),
    jnp.float32,
)

逻辑：

kv_block_index = pl.program_id(3)

@pl.when(kv_block_index == 0)
def init():
    dq_scratch[...] = 0

@pl.when(should_run)
def run():
    dq_scratch[...] += partial_dq

@pl.when(kv_block_index == last_kv_block)
def store():
    dq_ref[...] = dq_scratch.astype(dq_ref.dtype)

4.3 DKV scratch 模式

DKV 需要沿 Q 维累加：

dK = sum_over_Q dS^T @ Q
dV = sum_over_Q P^T @ dO

所以 DKV scratch 可以设计为：

dk_scratch = pltpu.VMEM(
    (block_b, block_kv_major_dkv, head_dim),
    jnp.float32,
)

dv_scratch = pltpu.VMEM(
    (block_b, block_kv_major_dkv, head_dim),
    jnp.float32,
)

逻辑：

q_block_index = pl.program_id(3)

@pl.when(q_block_index == 0)
def init():
    dk_scratch[...] = 0
    dv_scratch[...] = 0

@pl.when(should_run)
def run():
    dk_scratch[...] += partial_dk
    dv_scratch[...] += partial_dv

@pl.when(q_block_index == last_q_block)
def store():
    dk_ref[...] = dk_scratch.astype(dk_ref.dtype)
    dv_ref[...] = dv_scratch.astype(dv_ref.dtype)

4.4 使用 scratch 的注意事项

使用 scratch 跨 grid 维累加时，要注意：

1. reduction grid 维必须有明确执行顺序；
2. 不能在完全 parallel 的 grid 维上假设先后顺序；
3. 初始化、累加、写出条件必须严格正确；
4. scratch shape 不要过大；
5. scratch 通常用 fp32 保证累加精度；
6. 最终写出时再 cast 到输出 dtype。

如果依赖某个 grid 维顺序推进，需要在 TPU compiler params 中正确设置该维的语义，例如将 reduction 维设置为 "arbitrary"，而不是所有维都当成完全 parallel。

5. BlockSpec 设计规范

5.1 BlockSpec 应该 tile 化

推荐：

q_spec = pl.BlockSpec(
    (block_b, group_size, block_q_major, head_dim),
    q_index_map,
)

不推荐：

q_spec = pl.BlockSpec(
    (block_b, group_size, q_seq_len, head_dim),
    q_full_index_map,
)

因为后者意味着：

对 q_seq_len 这个维度几乎没有 tile 化。

这会让一个 program 看到完整 sequence，容易导致：

1. input window 过大；
2. prefetch 粒度过粗；
3. compiler 难以优化；
4. live range 变长；
5. VMEM/register 压力变大；
6. MXU pipeline 不稳定。

5.2 DQ BlockSpec 示例

def q_index_map(batch_index, kv_head_index, q_block_index, kv_block_index):
    return batch_index, kv_head_index, q_block_index, 0

def kv_index_map(batch_index, kv_head_index, q_block_index, kv_block_index):
    return batch_index, kv_head_index, kv_block_index, 0

q_spec = pl.BlockSpec(
    (block_b, group_size, block_q_major_dq, head_dim),
    q_index_map,
)

kv_spec = pl.BlockSpec(
    (block_b, 1, block_kv_major_dq, head_dim),
    kv_index_map,
)

lse_spec = pl.BlockSpec(
    (block_b, group_size, block_q_major_dq, MIN_BLOCK_SIZE),
    q_index_map,
)

5.3 DKV BlockSpec 示例

def q_index_map(batch_index, kv_head_index, kv_block_index, q_block_index):
    return batch_index, kv_head_index, q_block_index, 0

def kv_index_map(batch_index, kv_head_index, kv_block_index, q_block_index):
    return batch_index, kv_head_index, kv_block_index, 0

q_spec = pl.BlockSpec(
    (block_b, group_size, block_q_major_dkv, head_dim),
    q_index_map,
)

kv_spec = pl.BlockSpec(
    (block_b, 1, block_kv_major_dkv, head_dim),
    kv_index_map,
)

lse_spec = pl.BlockSpec(
    (block_b, group_size, block_q_major_dkv, MIN_BLOCK_SIZE),
    q_index_map,
)

5.4 Skipped block 的 safe index

如果 causal 下某些 block 完全无效，Pallas 仍可能需要 index_map 返回合法 tile。

可以使用 safe index：

def kv_index_map(batch_index, kv_head_index, q_block_index, kv_block_index):
    should_run = below_or_on_diag(
        q_block_index,
        block_q_major,
        kv_block_index,
        block_kv_major,
    )
    safe_kv_index = lax.select(should_run, kv_block_index, 0)
    return batch_index, kv_head_index, safe_kv_index, 0

这样可以避免无效或越界 prefetch。

6. Major / Minor Tile 分层

6.1 为什么要分 major/minor？

不要只用一层：

block_q
block_kv

更推荐两层：

major block:
  grid 级别 tile

minor block:
  program 内部 dot tile

好处：

1. grid 粒度和 dot 粒度分开调；
2. 更容易平衡并行度和数据复用；
3. 更容易控制 VMEM/register 压力；
4. 更容易找到 MXU 友好的 matmul shape；
5. 更容易做 block size sweep；
6. 更接近成熟 FlashAttention kernel 的写法。

6.2 DQ tile 参数

推荐：

block_q_major_dq
block_kv_major_dq
block_kv_dq

含义：

block_q_major_dq:
  DQ kernel 一个 q grid block 覆盖的 Q 长度

block_kv_major_dq:
  DQ kernel 一个 kv grid block 覆盖的 KV 长度

block_kv_dq:
  DQ kernel 内部每次 dot 的 KV minor tile

示例：

block_q_major_dq = 128
block_kv_major_dq = 256
block_kv_dq = 128

表示：

一个 DQ program 覆盖 128 个 Q token 和 256 个 KV token，
但内部每次只用 128 个 KV token 做一次 dot。

伪代码：

for kv_minor in range(0, block_kv_major_dq, block_kv_dq):
    k = k_tile[:, :, kv_minor : kv_minor + block_kv_dq, :]
    v = v_tile[:, :, kv_minor : kv_minor + block_kv_dq, :]
    scores = q @ k.T
    ...
    dq_scratch += ds @ k

6.3 DKV tile 参数

推荐：

block_q_major_dkv
block_kv_major_dkv
block_q_dkv
block_kv_dkv

含义：

block_q_major_dkv:
  DKV kernel 一个 q grid block 覆盖的 Q 长度

block_kv_major_dkv:
  DKV kernel 一个 kv grid block 覆盖的 KV 长度

block_q_dkv:
  DKV kernel 内部每次处理的 Q minor tile

block_kv_dkv:
  DKV kernel 内部每次处理的 KV minor tile

示例：

block_q_major_dkv = 128
block_kv_major_dkv = 256
block_q_dkv = 128
block_kv_dkv = 128

表示：

一个 DKV program 覆盖 128 个 Q token 和 256 个 KV token，
但内部把 KV 再拆成两个 128 的 minor tile。

伪代码：

for q_minor in range(0, block_q_major_dkv, block_q_dkv):
    for kv_minor in range(0, block_kv_major_dkv, block_kv_dkv):
        q = q_tile[:, :, q_minor : q_minor + block_q_dkv, :]
        k = k_tile[:, :, kv_minor : kv_minor + block_kv_dkv, :]
        v = v_tile[:, :, kv_minor : kv_minor + block_kv_dkv, :]

        scores = q @ k.T
        ...
        dk_scratch[kv_minor] += ds.T @ q
        dv_scratch[kv_minor] += p.T @ do

7. Causal Mask 优化规范

7.1 区分 block-level mask 和 element-level mask

Element-level causal mask：

scores = scores + jnp.where(mask, 0.0, DEFAULT_MASK_VALUE)

这是必要的，因为对角线附近的 partial block 仍需要逐元素 mask。

但对于完全在 causal 上三角的 block，应该 block-level skip：

如果 q_block_end < kv_block_start，
则整个 block 无效，不需要做 QK^T。

7.2 Block-level skip

推荐函数：

def below_or_on_diag(q_block_index, block_q, kv_block_index, block_kv):
    q_block_end = (q_block_index + 1) * block_q - 1
    kv_block_start = kv_block_index * block_kv
    return q_block_end >= kv_block_start

DQ 中：

should_run = below_or_on_diag(
    q_block_index,
    block_q_major_dq,
    kv_block_index,
    block_kv_major_dq,
)

DKV 中：

should_run = below_or_on_diag(
    q_block_index,
    block_q_major_dkv,
    kv_block_index,
    block_kv_major_dkv,
)

然后：

@pl.when(should_run)
def run():
    ...

7.3 为什么 block-level skip 很重要？

对于 causal prefill，理论上上三角大约一半 attention block 无效。

如果不 skip，而是：

scores = q @ k.T
scores = apply_causal_mask(scores)

则无效 block 仍然消耗：

1. QK^T matmul
2. exp
3. dp
4. ds
5. dQ/dK/dV partial accumulation

所以 block-level skip 可以直接降低 latency。

8. Attention Backward 常用生产级模板

8.1 DQ kernel 模板

数学：

dQ = sum_over_KV dS @ K

结构：

def dq_kernel(
    q_ref,
    k_ref,
    v_ref,
    o_ref,
    lse_ref,
    do_ref,
    dq_ref,
    dq_scratch_ref,
    *,
    sm_scale,
    block_kv_minor,
    kv_seq_len,
):
    q_block_index = pl.program_id(2)
    kv_block_index = pl.program_id(3)

    @pl.when(kv_block_index == 0)
    def init():
        dq_scratch_ref[...] = jnp.zeros_like(dq_scratch_ref)

    should_run = below_or_on_diag(
        q_block_index,
        block_q_major_dq,
        kv_block_index,
        block_kv_major_dq,
    )

    @pl.when(should_run)
    def run():
        q = q_ref[...]
        o = o_ref[...]
        do = do_ref[...]
        lse = lse_ref[...]

        di = jnp.sum(
            o.astype(jnp.float32) * do.astype(jnp.float32),
            axis=-1,
        )

        for kv_minor in range(0, block_kv_major_dq, block_kv_minor):
            k = k_ref[..., kv_minor : kv_minor + block_kv_minor, :]
            v = v_ref[..., kv_minor : kv_minor + block_kv_minor, :]

            scores = dot(q, k.T)
            scores *= sm_scale
            scores = apply_causal_mask_if_needed(scores)

            p = exp(scores - lse)
            dp = dot(do, v.T)
            ds = (dp - di[:, None]) * p
            ds *= sm_scale

            dq_scratch_ref[...] += dot(ds, k)

    @pl.when(kv_block_index == kv_seq_len // block_kv_major_dq - 1)
    def store():
        dq_ref[...] = dq_scratch_ref.astype(dq_ref.dtype)

8.2 DKV kernel 模板

数学：

dK = sum_over_Q dS^T @ Q
dV = sum_over_Q P^T @ dO

结构：

def dkv_kernel(
    q_ref,
    k_ref,
    v_ref,
    o_ref,
    lse_ref,
    do_ref,
    dk_ref,
    dv_ref,
    dk_scratch_ref,
    dv_scratch_ref,
    *,
    sm_scale,
    block_q_minor,
    block_kv_minor,
    q_seq_len,
):
    kv_block_index = pl.program_id(2)
    q_block_index = pl.program_id(3)

    @pl.when(q_block_index == 0)
    def init():
        dk_scratch_ref[...] = jnp.zeros_like(dk_scratch_ref)
        dv_scratch_ref[...] = jnp.zeros_like(dv_scratch_ref)

    should_run = below_or_on_diag(
        q_block_index,
        block_q_major_dkv,
        kv_block_index,
        block_kv_major_dkv,
    )

    @pl.when(should_run)
    def run():
        for q_minor in range(0, block_q_major_dkv, block_q_minor):
            q = q_ref[..., q_minor : q_minor + block_q_minor, :]
            o = o_ref[..., q_minor : q_minor + block_q_minor, :]
            do = do_ref[..., q_minor : q_minor + block_q_minor, :]
            lse = lse_ref[..., q_minor : q_minor + block_q_minor, :]

            di = jnp.sum(
                o.astype(jnp.float32) * do.astype(jnp.float32),
                axis=-1,
            )

            for kv_minor in range(0, block_kv_major_dkv, block_kv_minor):
                k = k_ref[..., kv_minor : kv_minor + block_kv_minor, :]
                v = v_ref[..., kv_minor : kv_minor + block_kv_minor, :]

                scores = dot(q, k.T)
                scores *= sm_scale
                scores = apply_causal_mask_if_needed(scores)

                p = exp(scores - lse)
                dp = dot(do, v.T)
                ds = (dp - di[:, None]) * p
                ds *= sm_scale

                dk_scratch_ref[..., kv_minor, :] += dot(ds.T, q)
                dv_scratch_ref[..., kv_minor, :] += dot(p.T, do)

    @pl.when(q_block_index == q_seq_len // block_q_major_dkv - 1)
    def store():
        dk_ref[...] = dk_scratch_ref.astype(dk_ref.dtype)
        dv_ref[...] = dv_scratch_ref.astype(dv_ref.dtype)

9. LSE / m / l 设计规范

9.1 LSE 是什么？

LSE 是：

logsumexp(scores)

也就是：

lse_i = log(sum_j exp(scores_i,j))

backward 中可以用它重建 softmax：

p = jnp.exp(scores - lse[:, None])

9.2 m/l 和 LSE 的关系

FlashAttention forward 有时保存：

m = row max
l = sum(exp(scores - m))

那么：

lse = m + log(l)

两种 backward 重建 P 的方式等价：

p = exp(scores - lse)

等价于：

p = exp(scores - m) / l

9.3 生产级建议

1. 如果已有 LSE，backward 用 exp(scores - lse) 最简单；
2. 如果 forward 是 streaming softmax，可保存 m/l；
3. residuals 应尽量比完整 P 小；
4. 不要保存完整 [B, H, Q, K] attention matrix；
5. LSE/m/l 通常用 fp32。

10. Dtype 与数值稳定性规范

10.1 Matmul accumulation

建议：

preferred_element_type=jnp.float32

尤其是：

QK^T
P @ V
dO @ V^T
dS @ K
dS^T @ Q
P^T @ dO

10.2 Softmax 相关

建议：

scores / p / ds / accumulators 尽量使用 fp32
最终输出再 cast 到 q/k/v dtype

常见写法：

scores = dot(..., preferred_element_type=jnp.float32)
p = jnp.exp(scores - lse)
ds = (dp - di[:, None]) * p
dq_acc = dq_acc.astype(jnp.float32)

10.3 mask value

常见写法：

DEFAULT_MASK_VALUE = -0.7 * float(jnp.finfo(jnp.float32).max)

避免直接使用 -inf 导致某些平台上的数值或编译问题。

11. GQA / MQA Kernel 设计规范

11.1 GQA 的核心

GQA 中：

num_q_heads > num_kv_heads
group_size = num_q_heads // num_kv_heads

每个 KV head 被多个 Q heads 共享。

因此：

dK/dV 需要聚合来自同一 group 内多个 Q heads 的梯度。

11.2 推荐 layout

对于一个 KV head：

Q tile:
  [group_size, block_q, head_dim]

K/V tile:
  [1, block_kv, head_dim]

可以 reshape：

[group_size, block_q, head_dim]
->
[group_size * block_q, head_dim]

这样一次 matmul 覆盖整个 Q-head group。

11.3 group_size 的风险

如果：

group_size * block_q

太大，则：

scores / p / dp / ds

会非常大。

例如：

group_size = 8
block_q = 128
block_kv = 256
scores = [1024, 256]

这可能导致 VMEM/register 压力过大。

生产级策略：

第一阶段:
  block_g = full group_size
  简化 dK/dV reduction

第二阶段:
  如果 group_size 很大且 VMEM 压力高，再引入 block_g

11.4 block_g 的代价

如果引入：

block_g < group_size

则 DKV 需要额外处理：

不同 group block 对同一个 dK/dV 的累加。

这可能需要：

1. 更复杂的 scratch；
2. 额外 reduction kernel；
3. atomic-like accumulation 设计；
4. 更复杂的 correctness 测试。

所以不建议第一版就引入 block_g。

12. 性能优化 Checklist

12.1 MXU 利用率低时优先检查

1. dot_general shape 是否太小？
2. dot_general shape 是否不规整？
3. block size 是否 128 对齐？
4. 是否有过长 program 内部 loop？
5. reduction 维是否应该放进 grid？
6. 是否用了 full-seq BlockSpec？
7. scores/p/dp/ds 是否太大？
8. VMEM/register 是否压力太高？
9. 是否有 spill 或 compiler warning？
10. causal 上三角是否被 matmul 计算了？
11. matmul 中间是否夹杂太多 elementwise？
12. grid program 数是否太少？
13. batch/head 并行度是否足够？
14. block_b 是否过大或过小？

12.2 Latency 高时优先检查

1. 是否没有 block-level causal skip？
2. 是否物化了 [B, H, Q, K] 大 tensor？
3. 是否写出了 dAB/dS？
4. 是否 HBM 读写过多？
5. 是否重复读取 full Q/K/V？
6. scratch 是否太大？
7. block_kv 是否太大导致临时矩阵膨胀？
8. block_q 是否太大导致 group_size * block_q 过大？
9. 是否 padding 后多算太多 token？
10. DQ 和 DKV 哪一个更慢？

12.3 Correctness 失败时优先检查

1. causal mask 的 row/col offset 是否正确？
2. q_start / kv_start 是否正确？
3. GQA head mapping 是否正确？
4. lse/m/l shape 是否正确 broadcast？
5. sm_scale 是否 forward/backward 一致？
6. ds 是否乘了 sm_scale？
7. di = sum(o * do) 是否按 row 计算？
8. dtype cast 是否过早？
9. padding 后是否正确 slice 回原 shape？
10. skipped block 是否误跳过了 diagonal block？
11. scratch init/store 条件是否正确？

13. Benchmark 规范

13.1 latency 是主指标

当比较 Pallas kernel 和 JAX reference 时，不要直接混用 FLOP 口径。

常见 FLOP 口径包括：

1. JAX compiler estimated FLOPs
2. manual Pallas matmul inventory
3. theoretical math FLOPs
4. actual executed FLOPs with causal skip

这些不是一回事。

速度比较以：

latency

为主。

13.2 effective TFLOP/s

如果要计算：

effective_TFLOP/s

必须固定 FLOP 口径。

建议使用：

manual Pallas matmul inventory

公式：

effective_TFLOP/s = manual_pallas_GFLOP / latency_ms

因为：

1 GFLOP / 1 ms = 1 TFLOP/s

13.3 必须分别测 DQ 和 DKV

不要只看 total。

建议记录：

DQ latency
DKV latency
Total latency
DQ MXU utilization
DKV MXU utilization
Total MXU utilization

原因：

DQ 和 DKV 的瓶颈可能完全不同。

例如：

DQ 可能卡在 KV reduction；
DKV 可能卡在 full Q BlockSpec、scratch、group reduction。

13.4 Benchmark 矩阵

建议至少测：

batch:
  1, 2, 4

seq_len:
  512, 1024, 2048, 4096, 8192

head_dim:
  64, 128

num_q_heads / num_kv_heads:
  16/4, 32/8, 32/4

group_size:
  2, 4, 8

dtype:
  bf16 input + fp32 accumulate

14. Block Size Sweep 规范

14.1 Forward 和 backward 的最优 block size 可能不同

不要假设：

forward 最优 block size = backward 最优 block size

原因：

forward:
  QK^T + P@V + streaming softmax

backward:
  重建 P
  dV = P^T dO
  dP = dO V^T
  dS
  dK = dS^T Q
  dQ = dS K

backward 的 matmul 数量、reduction 方向、scratch 需求都不同。

14.2 DQ 和 DKV 的最优 block size 也可能不同

DQ：

输出 dQ
沿 KV reduction

DKV：

输出 dK/dV
沿 Q reduction
还要聚合 group_size 个 Q heads

所以：

block_q_major_dq
block_kv_major_dq
block_kv_dq

和：

block_q_major_dkv
block_kv_major_dkv
block_q_dkv
block_kv_dkv

应该分开调。

14.3 推荐初始 sweep

DQ：

block_q_major_dq:
  64, 128

block_kv_major_dq:
  128, 256

block_kv_dq:
  128

DKV：

block_q_major_dkv:
  64, 128, 256

block_kv_major_dkv:
  128, 256

block_q_dkv:
  64, 128

block_kv_dkv:
  128

如果 group_size 较大：

block_q_major:
  优先尝试 64, 128

block_kv_major:
  优先尝试 128

15. 生产级代码结构建议

15.1 参数 dataclass

推荐：

@dataclasses.dataclass(frozen=True)
class BlockSizes:
    block_b: int = 1

    block_q_major_dq: int = 128
    block_kv_major_dq: int = 128
    block_kv_dq: int = 128

    block_q_major_dkv: int = 128
    block_kv_major_dkv: int = 128
    block_q_dkv: int = 128
    block_kv_dkv: int = 128

    def __post_init__(self):
        ...

校验：

1. 所有 block size > 0；
2. minor <= major；
3. major % minor == 0；
4. block_kv_* 是 128 的倍数；
5. q_seq_len / kv_seq_len 是否可整除，或是否需要 padding。

15.2 Shape validation

在外层 API 做 shape 校验：

def validate_shapes(q, k, v, o, lse, do):
    assert q.shape == o.shape
    assert q.shape == do.shape
    assert k.shape == v.shape
    assert q.shape[1] % k.shape[1] == 0
    assert q.shape[-1] == k.shape[-1]

不要把 shape 错误留到 kernel 内部才失败。

15.3 Padding 和 slicing

生产级 kernel 通常需要支持非整除 sequence length。

做法：

q_pad = pad_axis_to_multiple(q, axis=2, multiple=block_q)
k_pad = pad_axis_to_multiple(k, axis=2, multiple=block_kv)
...
out = out[:, :, :original_seq_len, :]

注意：

padding token 必须被 mask 或保证不影响结果。

15.4 named_scope

推荐：

name_scope = (
    f"gqa_bwd_dq_"
    f"{block_q_major_dq=}_"
    f"{block_kv_major_dq=}_"
    f"{block_kv_dq=}"
)

with jax.named_scope(name_scope):
    ...

好处：

1. profile 更容易看；
2. HLO / trace 更容易定位；
3. benchmark 结果更清楚。

15.5 debug / interpret

建议外层 API 支持：

debug: bool = False
interpret: bool = False

开发阶段：

interpret=True 方便调试

性能测试阶段：

interpret=False
debug=False

16. 常见反模式

16.1 一个 program 扫完整 seq_len

不推荐：

for q_start in range(0, q_seq_len, block_q):
    ...

如果这个 loop 很长，应考虑把 q_block 放入 grid。

16.2 full-sequence BlockSpec

不推荐：

pl.BlockSpec(
    (block_b, group_size, q_seq_len, head_dim),
    index_map,
)

推荐：

pl.BlockSpec(
    (block_b, group_size, block_q_major, head_dim),
    index_map,
)

16.3 先 matmul 再 mask 掉整个无效 block

不推荐：

scores = q @ k.T
scores = causal_mask(scores)

对于完全无效 block，应先 skip：

if block_is_valid:
    scores = q @ k.T

16.4 写出 dS / dAB 大矩阵

如果不是必须，不要写：

dAB = dS = [B, H, Q, K]

这会破坏 FlashAttention 不物化 attention matrix 的优势。

16.5 盲目增大 block_kv

block_kv=256 不一定比 128 快。

需要看：

1. MXU utilization
2. VMEM pressure
3. scores/p/dp/ds size
4. causal skip 粒度
5. latency

17. 推荐优化顺序

如果一个 Pallas attention backward kernel MXU utilization 很低，建议按下面顺序优化。

Step 1: 分离 DQ / DKV profiling

先确认：

DQ 慢，还是 DKV 慢？

记录：

DQ latency
DKV latency
DQ MXU utilization
DKV MXU utilization

Step 2: 加 block-level causal skip

这是低风险高收益优化。

目标：

完全 causal invalid 的 block 不做 QK^T。

Step 3: 消除 full-seq BlockSpec

将：

(block_b, group_size, q_seq_len, head_dim)

改成：

(block_b, group_size, block_q_major, head_dim)

Step 4: reduction 维放进 grid

DQ：

把 kv_block 放进 grid。

DKV：

把 q_block 放进 grid。

Step 5: 用 scratch 做 partial accumulation

DQ：

dq_scratch 跨 kv_block 累加。

DKV：

dk_scratch / dv_scratch 跨 q_block 累加。

Step 6: 引入 major/minor tile

把：

block_q
block_kv

拆成：

block_q_major
block_kv_major
block_q_minor
block_kv_minor

Step 7: 系统 sweep block size

不要凭感觉判断最优 block。

至少 sweep：

block_q_major = 64 / 128 / 256
block_kv_major = 128 / 256
block_q_minor = 64 / 128
block_kv_minor = 128

Step 8: 视情况引入 block_g

只有当：

group_size 很大
scores tile 太大
VMEM 压力明显

才考虑：

block_g < group_size

18. 代码评审 Checklist

提交生产级 Pallas kernel 前，建议检查：

[ ] 数学公式是否写清楚？
[ ] 输入输出 shape 是否校验？
[ ] GQA head mapping 是否正确？
[ ] grid 是否暴露足够并行度？
[ ] 是否避免 full-seq BlockSpec？
[ ] 是否避免单 program 内部长 reduction loop？
[ ] 是否正确使用 scratch？
[ ] scratch init/store 条件是否正确？
[ ] reduction grid 维是否有正确 dimension semantics？
[ ] causal block-level skip 是否正确？
[ ] diagonal block 是否仍然做 element-level mask？
[ ] dtype cast 是否合理？
[ ] accumulation 是否使用 fp32？
[ ] 是否避免写出不必要的大 tensor？
[ ] padding 后是否 slice 回原 shape？
[ ] 是否有 reference correctness test？
[ ] 是否有 DQ/DKV 单独 benchmark？
[ ] 是否记录 latency / MXU utilization / effective TFLOP/s？
[ ] 是否明确 FLOP 统计口径？
[ ] 是否 sweep 过核心 block size？
[ ] 是否有 named_scope 方便 profile？

19. 一句话总结

生产级 Pallas kernel 的核心不是“把公式翻译成 kernel”，而是：

把数学计算拆成硬件友好的 tile，
把长 reduction 暴露成 grid，
用 scratch 做局部累加，
避免 full-sequence input window，
尽量让 MXU 连续吃到规整 matmul，
同时控制 VMEM/register/HBM 压力。

对于 FlashAttention / GQA backward，最常用、最值得掌握的模式是：

DQ:
  grid over q_block × kv_block
  dq_scratch 跨 kv_block 累加

DKV:
  grid over kv_block × q_block
  dk_scratch / dv_scratch 跨 q_block 累加

Causal:
  block-level skip + diagonal block element-level mask

Tiling:
  major block 控制 grid 粒度
  minor block 控制 program 内部 dot 粒度

Benchmark:
  latency 为主
  FLOP 口径固定
  DQ/DKV 分开看
]]>Pallas TPU Distributed Collectives 复习文档2026-06-04T00:00:00+00:002026-06-04T00:00:00+00:00https://wqh011128.github.io/blog/2026/06/04/Pallas_TPU_Distributed_CollectivesPallas TPU Distributed Collectives 复习文档

本文根据 JAX 官方文档 Distributed Computing in Pallas for TPUs 整理，重点复习四个 collective 示例：

  • lax.ppermute
  • lax.all_gather
  • lax.psum
  • lax.psum_scatter

说明：代码是“教学整理版”，保留官方示例的核心结构、API 和数据流，但注释与组织方式做了重写，方便复习。真实运行仍需要 TPU 环境。

0. 共同背景

0.1 Ring 设备模型

假设有 4 个 device：

D0 -> D1 -> D2 -> D3 -> D0

右邻居：

right = (my_id + 1) % num_devices

左邻居：

left = (my_id - 1 + num_devices) % num_devices

在本文模拟中：

D0 的本地数据记作 x0
D1 的本地数据记作 x1
D2 的本地数据记作 x2
D3 的本地数据记作 x3

0.2 mesh、PartitionSpec、NamedSharding

P = jax.sharding.PartitionSpec

mesh = jax.make_mesh((num_devices,), ("x",))
partition = P(None, "x")
sharding = jax.sharding.NamedSharding(mesh, partition)

三者区别：

名称 含义
mesh 设备组成的逻辑网格，例如一维 x 轴有 4 个 TPU device
PartitionSpec 说明数组维度如何映射到 mesh 轴
NamedSharding 把设备网格和切分规则合成一个真正可用于 device_put 的 sharding

例子：

partition = P(None, "x")
input_arr.shape == (8, 128 * num_devices)

表示：

第 0 维不切分
第 1 维沿 x 轴切分

如果 num_devices=4，全局 shape 是 (8, 512)，每个 device 本地拿到 (8, 128)。

0.3 TPU Pallas 通信：RDMA 是 push-only

TPU Pallas 的 remote DMA 是 push-only：

当前 device 可以把本地 src_ref 推送到远端 device 的 dst_ref。
当前 device 不能直接读取远端 device 的内存。

核心 API：

copy = pltpu.make_async_remote_copy(
    src_ref=local_ref,
    dst_ref=remote_ref,
    send_sem=send_sem,
    recv_sem=recv_sem,
    device_id=(target_device,),
    device_id_type=pl.DeviceIdType.MESH,
)
copy.start()
copy.wait()

关键点：

参数 含义
src_ref 本 device 上要发送的数据
dst_ref 目标 device 上接收的位置
send_sem 发送端 DMA semaphore
recv_sem 接收端 DMA semaphore
.start() 发起异步 DMA
.wait() 等发送和接收都完成

0.4 HBM、VMEM、VREG 的直觉

在这些示例里，常见内存角色是：

名称 直觉
HBM 每个 device 自己的大内存，适合放通信 buffer
VMEM TPU 上更靠近计算的内存，适合做局部计算
VREG vector register，真正执行向量操作的位置

一个常见流程：

远端 DMA 写入本 device 的 HBM scratch
        ↓
local DMA 把 HBM scratch 拷到 VMEM scratch
        ↓
在 VMEM/VREG 中计算
        ↓
写回 HBM scratch 或最终 o_ref

1. lax.ppermute：右移一格

1.1 直观解释

ppermute 是最简单的跨设备通信：每个 device 把自己的 shard 发给右邻居。

对 4 个 device：

输入:
  D0: x0
  D1: x1
  D2: x2
  D3: x3

每个 device 发送到右邻居:
  D0 -> D1
  D1 -> D2
  D2 -> D3
  D3 -> D0

输出:
  D0: x3
  D1: x0
  D2: x1
  D3: x2

1.2 配图

flowchart LR
  D0["D0: x0"] --> D1["D1 receives x0"]
  D10["D1: x1"] --> D2["D2 receives x1"]
  D20["D2: x2"] --> D3["D3 receives x2"]
  D30["D3: x3"] --> D0R["D0 receives x3"]

1.3 教学版完整代码

import jax
from jax import lax
from jax import numpy as jnp
from jax.experimental import pallas as pl
from jax.experimental.pallas import tpu as pltpu

P = jax.sharding.PartitionSpec
num_devices = jax.local_device_count()

partition = P(None, "x")
mesh = jax.make_mesh((num_devices,), ("x",))
sharding = jax.sharding.NamedSharding(mesh, partition)

input_arr = jax.random.uniform(
    jax.random.key(0),
    shape=(8, 128 * num_devices),
)
input_arr = jax.device_put(input_arr, sharding)


def right_shift_kernel(x_ref, y_ref, send_sem, recv_sem):
    my_id = lax.axis_index("x")
    target = lax.rem(my_id + 1, num_devices)

    dma = pltpu.make_async_remote_copy(
        src_ref=x_ref,
        dst_ref=y_ref,
        send_sem=send_sem,
        recv_sem=recv_sem,
        device_id=(target,),
        device_id_type=pl.DeviceIdType.MESH,
    )
    dma.start()
    dma.wait()


out_shape = jax.ShapeDtypeStruct((8, 128), jnp.float32)

grid_spec = pltpu.PrefetchScalarGridSpec(
    num_scalar_prefetch=0,
    in_specs=[pl.BlockSpec(memory_space=pl.ANY)],
    out_specs=pl.BlockSpec(memory_space=pl.ANY),
    scratch_shapes=([pltpu.SemaphoreType.DMA] * 2),
)

right_shift = pl.pallas_call(
    right_shift_kernel,
    out_shape=out_shape,
    grid_spec=grid_spec,
)

pallas_result = jax.jit(
    jax.shard_map(
        right_shift,
        mesh=mesh,
        in_specs=partition,
        out_specs=partition,
        check_vma=False,
    )
)(input_arr)

perm = tuple((src, (src + 1) % num_devices) for src in range(num_devices))
xla_result = jax.jit(
    jax.shard_map(
        lambda x: lax.ppermute(x, "x", perm),
        mesh=mesh,
        in_specs=partition,
        out_specs=partition,
    )
)(input_arr)

1.4 逐步模拟

Device 初始本地输入 发送到 最终本地输出
D0 x0 D1 x3
D1 x1 D2 x0
D2 x2 D3 x1
D3 x3 D0 x2

完成这个操作的代码是：

target = (my_id + 1) % num_devices
make_async_remote_copy(src_ref=x_ref, dst_ref=y_ref, device_id=(target,))

这里 dst_ref=y_ref 表示目标 device 上同一个 Pallas output buffer 的本地位置。

2. lax.all_gather：每个 device 收集所有 shard

2.1 直观解释

all_gather 的目标是：

每个 device 一开始只有自己的 shard。
最后每个 device 都有所有 shard。

4 个 device：

初始:
  D0: x0
  D1: x1
  D2: x2
  D3: x3

最终每个 device:
  [x0, x1, x2, x3]

官方图：

all_gather

2.2 配图

flowchart LR
  S0["step 0: 每个 device 发自己的 shard"] --> S1["step 1: 发刚从左边收到的 shard"]
  S1 --> S2["step 2: 继续转发"]
  S2 --> OUT["每个 device 都拥有 x0,x1,x2,x3"]

2.3 教学版完整代码

partition = P("x", None)
mesh = jax.make_mesh((num_devices,), ("x",))
sharding = jax.sharding.NamedSharding(mesh, partition)

input_arr = jax.random.uniform(
    jax.random.key(0),
    shape=(8 * num_devices, 128),
)
input_arr = jax.device_put(input_arr, sharding)


def all_gather_kernel(x_ref, out_ref, local_sem, send_sem, recv_sems):
    step = pl.program_id(0)
    my_id = lax.axis_index("x")
    right = lax.rem(my_id + 1, num_devices)

    slot = lax.rem(my_id - step + num_devices, num_devices)

    @pl.when(step == 0)
    def _():
        local = pltpu.make_async_copy(
            src_ref=x_ref,
            dst_ref=out_ref.at[my_id],
            sem=local_sem,
        )
        local.start()
        local.wait()

    remote = pltpu.make_async_remote_copy(
        src_ref=out_ref.at[slot],
        dst_ref=out_ref.at[slot],
        send_sem=send_sem,
        recv_sem=recv_sems.at[step],
        device_id=(right,),
        device_id_type=pl.DeviceIdType.MESH,
    )
    remote.start()
    remote.wait()


out_shape = jax.ShapeDtypeStruct((num_devices, 8, 128), jnp.float32)

grid_spec = pltpu.PrefetchScalarGridSpec(
    num_scalar_prefetch=0,
    in_specs=[pl.BlockSpec(memory_space=pl.ANY)],
    out_specs=pl.BlockSpec(memory_space=pl.ANY),
    scratch_shapes=(
        [pltpu.SemaphoreType.DMA] * 2
        + [pltpu.SemaphoreType.DMA((num_devices - 1,))]
    ),
    grid=(num_devices - 1,),
)

all_gather = pl.pallas_call(
    all_gather_kernel,
    out_shape=out_shape,
    grid_spec=grid_spec,
)

pallas_result = jax.jit(
    jax.shard_map(
        all_gather,
        mesh=mesh,
        in_specs=partition,
        out_specs=partition,
        check_vma=False,
    )
)(input_arr)

xla_result = jax.jit(
    jax.shard_map(
        lambda x: lax.all_gather(x, "x"),
        mesh=mesh,
        in_specs=partition,
        out_specs=partition,
    )
)(input_arr)

2.4 关键代码解释

grid=(num_devices - 1,)

每个 device 自己已经有自己的 shard，只需要再收到其他 num_devices - 1 个 shard。

slot = (my_id - step) % num_devices

当前这一轮要把哪个 slot 转发给右邻居。

例如 D0：

step 0: 发送 slot 0
step 1: 发送 slot 3
step 2: 发送 slot 2

同时 D0 会从 D3 收到：

step 0: slot 3
step 1: slot 2
step 2: slot 1

2.5 遍历模拟

表格展示每轮结束后，每个 device 的 output slots。

loop 后 D0 slots D1 slots D2 slots D3 slots
初始 [-,-,-,-] [-,-,-,-] [-,-,-,-] [-,-,-,-]
step 0 [x0,-,-,x3] [x0,x1,-,-] [-,x1,x2,-] [-,-,x2,x3]
step 1 [x0,-,x2,x3] [x0,x1,-,x3] [x0,x1,x2,-] [-,x1,x2,x3]
step 2 [x0,x1,x2,x3] [x0,x1,x2,x3] [x0,x1,x2,x3] [x0,x1,x2,x3]

操作来源：

动作 代码
step 0 本地写自己的 slot out_ref.at[my_id] = x_ref via make_async_copy
每轮转发已有 slot src_ref=out_ref.at[slot]
写到右邻居同编号 slot dst_ref=out_ref.at[slot] with device_id=(right,)

为什么 recv_sems 有多个？

recv_sems.at[step]

每一轮一个 receive semaphore，避免多个 in-flight DMA 复用同一个接收计数器。官方文档也提醒，大 kernel 中 device 更容易不同步，复用 receive semaphore 可能导致静默错误。

3. lax.psum：单向 ring all-reduce sum

3.1 直观解释

psum 是 all-reduce sum：

每个 device 一开始有自己的 x。
最后每个 device 都得到 x0 + x1 + x2 + x3。

这个 Pallas 示例不是先 all-gather 再求和，而是：

数据沿 ring 传递。
每个 device 收到一份数据就加进自己的 accumulator。

官方图：

reduce_sum_2

3.2 配图

flowchart LR
  D0["D0: x0"] --> D1["D1 accum"]
  D1 --> D2["D2 accum"]
  D2 --> D3["D3 accum"]
  D3 --> D0R["D0 accum"]

每个 device 都在做同样的事情，所以一圈后每个 device 都累加了所有输入。

3.3 教学版完整代码

partition = P(None, "x")
mesh = jax.make_mesh((num_devices,), ("x",))
sharding = jax.sharding.NamedSharding(mesh, partition)

input_arr = jax.random.uniform(
    jax.random.key(0),
    shape=(8, 128 * num_devices),
)
input_arr = jax.device_put(input_arr, sharding)


def neighbor_barrier(left, right, double_barrier=True):
    barrier = pltpu.get_barrier_semaphore()
    for nbr in [left, right]:
        pl.semaphore_signal(
            barrier,
            inc=1,
            device_id=(nbr,),
            device_id_type=pl.DeviceIdType.MESH,
        )
    pl.semaphore_wait(barrier, 2)

    if double_barrier:
        @functools.partial(pl.run_scoped, second=pltpu.SemaphoreType.REGULAR)
        def _(second):
            for nbr in [left, right]:
                pl.semaphore_signal(
                    second,
                    inc=1,
                    device_id=(nbr,),
                    device_id_type=pl.DeviceIdType.MESH,
                )
            pl.semaphore_wait(second, 2)


def psum_kernel(
    x_ref,
    o_ref,
    hbm_buf,
    local_copy_sem,
    remote_recv_sem,
    remote_send_sem,
    capacity_sem,
    vmem_recv,
):
    step = pl.program_id(0)
    working = lax.rem(step, 2)
    receiving = 1 - working

    my_id = lax.axis_index("x")
    right = lax.rem(my_id + 1, num_devices)
    left = lax.rem(my_id - 1 + num_devices, num_devices)

    @pl.when(step == 0)
    def _():
        neighbor_barrier(left, right)
        o_ref[...] = jnp.zeros_like(o_ref)
        vmem_recv[...] = jnp.zeros_like(vmem_recv)

        first = pltpu.make_async_remote_copy(
            src_ref=x_ref,
            dst_ref=hbm_buf.at[working],
            send_sem=remote_send_sem,
            recv_sem=remote_recv_sem,
            device_id=(right,),
            device_id_type=pl.DeviceIdType.MESH,
        )
        first.start()
        first.wait()

    # 告诉左邻居：我已经准备好接收它下一次写入。
    pl.semaphore_signal(
        capacity_sem,
        inc=1,
        device_id=(left,),
        device_id_type=pl.DeviceIdType.MESH,
    )

    # 本地 HBM -> VMEM，用于累加。
    local = pltpu.make_async_copy(
        src_ref=hbm_buf.at[working],
        dst_ref=vmem_recv,
        sem=local_copy_sem,
    )
    local.start()

    # 写右邻居前，先等右邻居说它准备好了。
    pl.semaphore_wait(capacity_sem, 1)

    remote = pltpu.make_async_remote_copy(
        src_ref=hbm_buf.at[working],
        dst_ref=hbm_buf.at[receiving],
        send_sem=remote_send_sem,
        recv_sem=remote_recv_sem,
        device_id=(right,),
        device_id_type=pl.DeviceIdType.MESH,
    )
    remote.start()

    local.wait()
    o_ref[...] += vmem_recv[...]
    remote.wait()


out_shape = (
    jax.ShapeDtypeStruct((8, 128), jnp.float32),
    jax.ShapeDtypeStruct((2, 8, 128), jnp.float32),
)

grid_spec = pltpu.PrefetchScalarGridSpec(
    num_scalar_prefetch=0,
    in_specs=[pl.BlockSpec(memory_space=pltpu.VMEM)],
    out_specs=[
        pl.BlockSpec(memory_space=pltpu.VMEM),
        pl.BlockSpec(memory_space=pl.ANY),
    ],
    grid=(num_devices,),
    scratch_shapes=(
        [pltpu.SemaphoreType.DMA] * 3
        + [pltpu.SemaphoreType.REGULAR]
        + [pltpu.VMEM((8, 128), jnp.float32)]
    ),
)

psum_pallas = pl.pallas_call(
    psum_kernel,
    out_shape=out_shape,
    grid_spec=grid_spec,
    compiler_params=pltpu.CompilerParams(collective_id=0),
)

pallas_result = jax.jit(
    jax.shard_map(
        psum_pallas,
        mesh=mesh,
        in_specs=partition,
        out_specs=partition,
        check_vma=False,
    )
)(input_arr)[0]

xla_result = jax.jit(
    jax.shard_map(
        lambda x: lax.psum(x, "x"),
        mesh=mesh,
        in_specs=partition,
        out_specs=partition,
    )
)(input_arr)

3.4 关键区域

区域 memory space 用途
x_ref VMEM 本 device 的原始输入
o_ref VMEM 本 device 的累加结果
hbm_buf[0] HBM 通信双缓冲 slot 0
hbm_buf[1] HBM 通信双缓冲 slot 1
vmem_recv VMEM 从 HBM 拷出来后用于累加的临时 buffer

3.5 遍历模拟

H0 = hbm_buf[0]，H1 = hbm_buf[1]。

loop 后 D0 (H0,H1,o) D1 (H0,H1,o) D2 (H0,H1,o) D3 (H0,H1,o)
prologue 后 (x3,-,0) (x0,-,0) (x1,-,0) (x2,-,0)
step 0 (x3,x2,x3) (x0,x3,x0) (x1,x0,x1) (x2,x1,x2)
step 1 (x1,x2,x3+x2) (x2,x3,x0+x3) (x3,x0,x1+x0) (x0,x1,x2+x1)
step 2 (x1,x0,x3+x2+x1) (x2,x1,x0+x3+x2) (x3,x2,x1+x0+x3) (x0,x3,x2+x1+x0)
step 3 (x3,x0,all) (x0,x1,all) (x1,x2,all) (x2,x3,all)

其中：

all = x0 + x1 + x2 + x3

每轮核心操作来源：

动作 代码
第 0 轮先发自己的输入给右邻居 first = make_async_remote_copy(src_ref=x_ref, dst_ref=hbm_buf.at[working])
读当前 working slot 到 VMEM make_async_copy(src_ref=hbm_buf.at[working], dst_ref=vmem_recv)
把当前 working slot 继续发给右邻居 make_async_remote_copy(src_ref=hbm_buf.at[working], dst_ref=hbm_buf.at[receiving])
累加 o_ref[...] += vmem_recv[...]
防止跑快一轮覆盖邻居 working slot capacity_sem signal/wait

capacity_sem 的作用：

如果所有 device 严格同步，双缓冲不会冲突。
但不同 device 可以跑快或跑慢。
快的 device 下一轮 receiving_slot 可能正好是慢邻居还在读的 working_slot。
capacity_sem 让发送者必须等接收者确认“我已经进入对应轮次，可以写我的 receiving slot”。

4. lax.psum_scatter：双向 reduce-scatter

4.1 直观解释

psum_scatter 的语义可以看成：

先对所有 device 的输入按 block 求和；
然后每个 device 只保留属于自己的那个 block。

但高效实现不是：

all-reduce 完整结果 -> 再 scatter

而是：

每个 output block 的 partial sum 在 ring 上移动；
经过一个 device，就加上该 device 对这个 block 的贡献；
最后回到目标 device 时，这个 block 已经完整。

官方语义图：

reduce_scatter_1

官方通信图：

reduce_scatter_2

4.2 配图

一个 block 被切成上下两半：

T = top half    -> 向左传
B = bottom half -> 向右传

flowchart LR
  subgraph Top["top half: 向左"]
    T0["D0 starts T0"] --> T3["D3 adds"] --> T2["D2 adds"] --> T1["D1 adds"] --> T0R["D0 receives full T0"]
  end
  subgraph Bottom["bottom half: 向右"]
    B0["D0 starts B0"] --> B1["D1 adds"] --> B2["D2 adds"] --> B3["D3 adds"] --> B0R["D0 receives full B0"]
  end

注意：这段教学代码的数据流是双向的，但 phase 调度是交替的：

phase LEFT:
  发送上一阶段的 right-half
  计算当前 left-half

phase RIGHT:
  发送刚算好的 left-half
  计算当前 right-half

4.3 教学版完整代码

partition = P(None, "x")
mesh = jax.make_mesh((num_devices,), ("x",))
sharding = jax.sharding.NamedSharding(mesh, partition)

block_size = (16, 128)
input_arr = jax.random.uniform(
    jax.random.key(0),
    shape=(block_size[0] * num_devices, block_size[1] * num_devices),
)
input_arr = jax.device_put(input_arr, sharding)

LEFT = 0
RIGHT = 1


def mod(x, n):
    return lax.rem(x + n, n)


def signal(direction, sem):
    my_id = lax.axis_index("x")
    if direction == LEFT:
        target = mod(my_id - 1, num_devices)
    else:
        target = mod(my_id + 1, num_devices)
    pl.semaphore_signal(
        sem,
        inc=1,
        device_id=(target,),
        device_id_type=pl.DeviceIdType.MESH,
    )


def reduce_scatter_kernel(
    x_ref,
    o_ref,
    hbm_buf,
    local_copy_sem,
    left_recv_sem,
    left_send_sem,
    right_recv_sem,
    right_send_sem,
    left_capacity_sem,
    right_capacity_sem,
    accum,
):
    step = pl.program_id(0)
    phase = pl.program_id(1)

    is_first = jnp.logical_and(step == 0, phase == LEFT)
    is_last_step = step == pl.num_programs(0) - 1

    working = lax.rem(step, 2)
    receiving = 1 - working

    my_id = lax.axis_index("x")
    right = mod(my_id + 1, num_devices)
    left = mod(my_id - 1, num_devices)

    left_block = mod(my_id + step + 1, num_devices)
    right_block = mod(my_id - step - 1, num_devices)

    half_rows = block_size[0] // 2
    top = pl.ds(0, half_rows)
    bottom = pl.ds(half_rows, half_rows)
    current_half = pl.ds(phase * half_rows, half_rows)

    init_left = pltpu.make_async_remote_copy(
        src_ref=x_ref.at[my_id, top],
        dst_ref=hbm_buf.at[working, top],
        send_sem=left_send_sem,
        recv_sem=left_recv_sem,
        device_id=(left,),
        device_id_type=pl.DeviceIdType.MESH,
    )

    init_right = pltpu.make_async_remote_copy(
        src_ref=x_ref.at[my_id, bottom],
        dst_ref=hbm_buf.at[working, bottom],
        send_sem=right_send_sem,
        recv_sem=right_recv_sem,
        device_id=(right,),
        device_id_type=pl.DeviceIdType.MESH,
    )

    send_left = pltpu.make_async_remote_copy(
        src_ref=hbm_buf.at[working, top],
        dst_ref=hbm_buf.at[receiving, top],
        send_sem=left_send_sem,
        recv_sem=left_recv_sem,
        device_id=(left,),
        device_id_type=pl.DeviceIdType.MESH,
    )

    send_right = pltpu.make_async_remote_copy(
        src_ref=hbm_buf.at[receiving, bottom],
        dst_ref=hbm_buf.at[working, bottom],
        send_sem=right_send_sem,
        recv_sem=right_recv_sem,
        device_id=(right,),
        device_id_type=pl.DeviceIdType.MESH,
    )

    @pl.when(is_first)
    def _():
        neighbor_barrier(left, right)
        o_ref[...] = jnp.zeros_like(o_ref)
        accum[...] = jnp.zeros_like(accum)

        init_left.start()
        init_left.wait()
        init_right.start()

        signal(LEFT, right_capacity_sem)
        signal(RIGHT, left_capacity_sem)

    @pl.when(~is_first)
    def _():
        @pl.when(phase == LEFT)
        def _():
            pl.semaphore_wait(right_capacity_sem, 1)
            send_right.start()

        @pl.when(phase == RIGHT)
        def _():
            pl.semaphore_wait(left_capacity_sem, 1)
            send_left.start()

    local = pltpu.make_async_copy(
        src_ref=hbm_buf.at[working, current_half],
        dst_ref=accum,
        sem=local_copy_sem,
    )
    local.start()
    local.wait()

    @pl.when(~is_last_step)
    def _():
        @pl.when(phase == LEFT)
        def _():
            accum[...] += x_ref[left_block, top]

        @pl.when(phase == RIGHT)
        def _():
            accum[...] += x_ref[right_block, bottom]

    local = pltpu.make_async_copy(
        src_ref=accum,
        dst_ref=hbm_buf.at[working, current_half],
        sem=local_copy_sem,
    )
    local.start()
    local.wait()

    @pl.when(is_first)
    def _():
        init_right.wait()

    @pl.when(~is_first)
    def _():
        @pl.when(phase == LEFT)
        def _():
            send_right.wait()
            signal(LEFT, right_capacity_sem)

        @pl.when(phase == RIGHT)
        def _():
            send_left.wait()
            signal(RIGHT, left_capacity_sem)

    @pl.when(is_last_step)
    def _():
        @pl.when(phase == LEFT)
        def _():
            o_ref[top, ...] = accum[...]
            pl.semaphore_wait(right_capacity_sem, 1)

        @pl.when(phase == RIGHT)
        def _():
            o_ref[bottom, ...] = accum[...]
            pl.semaphore_wait(left_capacity_sem, 1)


out_shape = (
    jax.ShapeDtypeStruct((block_size[0], block_size[1]), jnp.float32),
    jax.ShapeDtypeStruct((2, block_size[0], block_size[1]), jnp.float32),
)

grid_spec = pltpu.PrefetchScalarGridSpec(
    num_scalar_prefetch=0,
    in_specs=[pl.BlockSpec(memory_space=pltpu.VMEM)],
    out_specs=[
        pl.BlockSpec(memory_space=pltpu.VMEM),
        pl.BlockSpec(memory_space=pl.ANY),
    ],
    grid=(num_devices, 2),
    scratch_shapes=(
        [pltpu.SemaphoreType.DMA] * 5
        + [pltpu.SemaphoreType.REGULAR] * 2
        + [pltpu.VMEM((block_size[0] // 2, block_size[1]), jnp.float32)]
    ),
)


def pallas_reduce_scatter(x):
    x = x.reshape(num_devices, block_size[0], block_size[1])
    return pl.pallas_call(
        reduce_scatter_kernel,
        out_shape=out_shape,
        grid_spec=grid_spec,
        compiler_params=pltpu.CompilerParams(collective_id=0),
    )(x)[0]


pallas_result = jax.jit(
    jax.shard_map(
        pallas_reduce_scatter,
        mesh=mesh,
        in_specs=P(None, "x"),
        out_specs=P("x", None),
        check_vma=False,
    )
)(input_arr)


def xla_reduce_scatter(x):
    x = x.reshape(num_devices, block_size[0], block_size[1])
    return lax.psum_scatter(x, "x")


xla_result = jax.jit(
    jax.shard_map(
        xla_reduce_scatter,
        mesh=mesh,
        in_specs=P(None, "x"),
        out_specs=P("x", None),
    )
)(input_arr)

4.4 输入为什么是设备倍数

input_arr.shape = (
    block_size[0] * num_devices,
    block_size[1] * num_devices,
)
partition = P(None, "x")

第 1 维被切到所有 device，所以每个 device 本地宽度是 block_size[1]。

每个 device 本地 shape 是：

(block_size[0] * num_devices, block_size[1])

然后 reshape：

x = x.reshape(num_devices, block_size[0], block_size[1])

本地变成：

x_d[0], x_d[1], x_d[2], x_d[3]

其中：

x_d[b] = device d 对最终 block b 的贡献

最终：

D0 得到 sum_d x_d[0]
D1 得到 sum_d x_d[1]
D2 得到 sum_d x_d[2]
D3 得到 sum_d x_d[3]

4.5 关键区域

区域 memory space 用途
x_ref VMEM 当前 device 的所有 block 贡献
o_ref VMEM 当前 device 最终保留的 reduced block
hbm_buf[0] HBM 通信 slot 0
hbm_buf[1] HBM 通信 slot 1
accum VMEM 当前 half-block 的临时累加器

4.6 遍历模拟

假设：

num_devices = 4
D0, D1, D2, D3
T = top half
B = bottom half

记号：

T2(2+1+0) = block 2 的 top half，已经累加 device 2、1、0 的贡献
B0(0+1+2+3) = block 0 的 bottom half 完整结果

Loop (0, LEFT)

执行：

init_left.start()
init_left.wait()
init_right.start()

local_copy: H0.T -> accum
accum += x_ref[left_block, T]
accum -> H0.T
Device H0.T H0.B H1.T H1.B accum o_ref
D0 T1(1+0) B3(3) - - T1(1+0) -
D1 T2(2+1) B0(0) - - T2(2+1) -
D2 T3(3+2) B1(1) - - T3(3+2) -
D3 T0(0+3) B2(2) - - T0(0+3) -

Loop (0, RIGHT)

执行：

send_left.start()
local_copy: H0.B -> accum
accum += x_ref[right_block, B]
accum -> H0.B
Device H0.T H0.B H1.T H1.B accum o_ref
D0 T1(1+0) B3(3+0) T2(2+1) - B3(3+0) -
D1 T2(2+1) B0(0+1) T3(3+2) - B0(0+1) -
D2 T3(3+2) B1(1+2) T0(0+3) - B1(1+2) -
D3 T0(0+3) B2(2+3) T1(1+0) - B2(2+3) -

Loop (1, LEFT)

执行：

send_right.start()
local_copy: H1.T -> accum
accum += x_ref[left_block, T]
accum -> H1.T
Device H0.T H0.B H1.T H1.B accum o_ref
D0 T1(1+0) B3(3+0) T2(2+1+0) B2(2+3) T2(2+1+0) -
D1 T2(2+1) B0(0+1) T3(3+2+1) B3(3+0) T3(3+2+1) -
D2 T3(3+2) B1(1+2) T0(0+3+2) B0(0+1) T0(0+3+2) -
D3 T0(0+3) B2(2+3) T1(1+0+3) B1(1+2) T1(1+0+3) -

Loop (1, RIGHT)

执行：

send_left.start()
local_copy: H1.B -> accum
accum += x_ref[right_block, B]
accum -> H1.B
Device H0.T H0.B H1.T H1.B accum o_ref
D0 T3(3+2+1) B3(3+0) T2(2+1+0) B2(2+3+0) B2(2+3+0) -
D1 T0(0+3+2) B0(0+1) T3(3+2+1) B3(3+0+1) B3(3+0+1) -
D2 T1(1+0+3) B1(1+2) T0(0+3+2) B0(0+1+2) B0(0+1+2) -
D3 T2(2+1+0) B2(2+3) T1(1+0+3) B1(1+2+3) B1(1+2+3) -

Loop (2, LEFT)

Device H0.T H0.B H1.T H1.B accum o_ref
D0 T3(3+2+1+0) B1(1+2+3) T2(2+1+0) B2(2+3+0) T3(3+2+1+0) -
D1 T0(0+3+2+1) B2(2+3+0) T3(3+2+1) B3(3+0+1) T0(0+3+2+1) -
D2 T1(1+0+3+2) B3(3+0+1) T0(0+3+2) B0(0+1+2) T1(1+0+3+2) -
D3 T2(2+1+0+3) B0(0+1+2) T1(1+0+3) B1(1+2+3) T2(2+1+0+3) -

Loop (2, RIGHT)

Device H0.T H0.B H1.T H1.B accum o_ref
D0 T3(3+2+1+0) B1(1+2+3+0) T0(0+3+2+1) B2(2+3+0) B1(1+2+3+0) -
D1 T0(0+3+2+1) B2(2+3+0+1) T1(1+0+3+2) B3(3+0+1) B2(2+3+0+1) -
D2 T1(1+0+3+2) B3(3+0+1+2) T2(2+1+0+3) B0(0+1+2) B3(3+0+1+2) -
D3 T2(2+1+0+3) B0(0+1+2+3) T3(3+2+1+0) B1(1+2+3) B0(0+1+2+3) -

Loop (3, LEFT)

最后一轮不再加本地贡献，只把完整 top half 写入输出。

Device H0.T H0.B H1.T H1.B accum o_ref
D0 T3(3+2+1+0) B1(1+2+3+0) T0(0+3+2+1) B0(0+1+2+3) T0(0+3+2+1) T0(all)
D1 T0(0+3+2+1) B2(2+3+0+1) T1(1+0+3+2) B1(1+2+3+0) T1(1+0+3+2) T1(all)
D2 T1(1+0+3+2) B3(3+0+1+2) T2(2+1+0+3) B2(2+3+0+1) T2(2+1+0+3) T2(all)
D3 T2(2+1+0+3) B0(0+1+2+3) T3(3+2+1+0) B3(3+0+1+2) T3(3+2+1+0) T3(all)

Loop (3, RIGHT)

最后一轮不再加本地贡献，只把完整 bottom half 写入输出。

Device H0.T H0.B H1.T H1.B accum o_ref
D0 T1(1+0+3+2) B1(1+2+3+0) T0(0+3+2+1) B0(0+1+2+3) B0(0+1+2+3) T0(all)+B0(all)
D1 T2(2+1+0+3) B2(2+3+0+1) T1(1+0+3+2) B1(1+2+3+0) B1(1+2+3+0) T1(all)+B1(all)
D2 T3(3+2+1+0) B3(3+0+1+2) T2(2+1+0+3) B2(2+3+0+1) B2(2+3+0+1) T2(all)+B2(all)
D3 T0(0+3+2+1) B0(0+1+2+3) T3(3+2+1+0) B3(3+0+1+2) B3(3+0+1+2) T3(all)+B3(all)

4.7 最后复习口诀

ppermute:
  一跳。每个 device 把自己的 shard 发给右邻居。

all_gather:
  多跳收集。每轮转发一个已有 slot，最终每个 device 收齐所有 slots。

psum:
  单向 all-reduce。数据沿 ring 转，每个 device 收到就加到自己的 o_ref。

psum_scatter:
  直接 reduce-scatter，不是真的先 all-reduce 再 scatter。
  partial sum 在路上传，经过 device 就加本地贡献。
  top half 向左，bottom half 向右。

5. 复习时最容易错的点

5.1 x_ref[i] 不是读 device i 的内存

在 shard_map 内，x_ref 永远是当前 device 的本地 shard。

所以：

x_ref[left_block, top]

意思是：

当前 device 本地保存的 block left_block 的 top half

不是远端 device left_block 的内存。

5.2 grid 循环不是天然全局同步

Pallas TPU grid 可以理解成本 device 上按顺序执行：

for step in range(...):
    kernel(...)

但不同 device 之间不保证每一轮严格同步。同步需要：

DMA semaphores
regular semaphores
barrier semaphores

5.3 双缓冲只解决同轮读写冲突

如果所有设备同轮执行：

working_slot 和 receiving_slot 不冲突

但如果一个 device 跑快一轮：

快 device 以为它在写邻居 receiving_slot
慢邻居可能还在读同一个 slot 作为 working_slot

所以 capacity_sem 用来做接收方确认。

5.4 psum_scatter 传的是 accumulator，不是输入本身

在 psum 中：

传输入 shard，accumulator 留在本地。

在 psum_scatter 中：

传 partial sum，输入贡献留在本地。

这是理解 reduce-scatter 的关键。

6. 官方参考

]]>JAX 到 VLIW，以及 Pallas / Splash Attention 复习文档2026-06-04T00:00:00+00:002026-06-04T00:00:00+00:00https://wqh011128.github.io/blog/2026/06/04/JAX_to_VLIW_Pallas_Splash_Attention_%E5%A4%8D%E4%B9%A0%E6%96%87%E6%A1%A3JAX 到 VLIW，以及 Pallas / Splash Attention 复习笔记

本文整理两篇 Patrick Toulme 的文章：

目标不是逐字翻译，而是按模型框架开发者的视角复习：重点理解 JAX/HLO/XLA/Pallas 这些上层如何影响底层执行；LLO/VLIW 只学会基础读法，不陷入指令细节。

0. 两条总线路：JAX to VLIW 与 Pallas to VLIW

先把两篇文章压缩成两张图。

0.1 普通 JAX / XLA 路线

Python JAX function
  -> JAX tracing / Jaxpr
  -> StableHLO / HLO
  -> HLO optimization passes
       algebraic simplification
       layout assignment
       tiling
       fusion
       memory assignment
       copy scheduling
  -> TPU backend LLO
  -> VLIW bundles
  -> TPU hardware
       HBM / VMEM / MXU / VPU / XLU / DMA

直觉：

JAX 代码描述“我要算什么”
HLO 描述“张量图是什么”
XLA 优化“怎样重排、融合、分配内存”
LLO 描述“TPU 上具体用哪些硬件动作”
VLIW bundle 描述“同一个时刻哪些硬件指令一起发射”

0.2 Pallas / Mosaic 路线

Python Pallas kernel
  -> pallas_call
  -> HLO custom-call
       XLA 只看到一个 opaque call
  -> Mosaic MLIR
       Mosaic 能看到 Pallas kernel body
  -> TPU backend LLO
  -> VLIW bundles
  -> TPU hardware

直觉：

普通 JAX:
  你写完整 tensor program，让 XLA 自动找 fusion / tiling。

Pallas:
  你直接写 tiled kernel program，让 Mosaic 和 TPU backend 继续降到底层。

0.3 两篇文章合起来的核心

第一篇的主题：

XLA 很强。
普通 JAX 程序可以被自动优化成 TPU 上的高质量 VLIW 程序。

第二篇的主题：

XLA 有边界。
XLA 会优化你写出来的算法，但通常不会自动发明另一个算法。
Splash / FlashAttention 这类 online softmax 是算法级改写，因此需要 Pallas 表达。

最重要的对比：

XLA fusion:
  减少已经存在的中间 tensor 的读写。

Splash / FlashAttention:
  改写算法，让完整 attention matrix 根本不成为程序语义里的 tensor。

第一篇：From JAX to VLIW

1. 例子在算什么

文章使用一个很小的 attention-like block，便于观察整个 TPU 编译链路：

h = x @ w1
h = h / sqrt(mean(h ** 2) + eps)
h = softmax(h)
out = h @ w2

形状大致是：

x:   [16, 64]
w1:  [64, 64]
w2:  [64, 32]
out: [16, 32]

它包含：

matmul_1
RMSNorm-like normalization
softmax
matmul_2

作者使用 jax.jit 触发编译，用 dump flags 输出 HLO 和 TPU backend 的 LLO。jax.named_call 的作用是让 IR metadata 中保留人类可读名字，例如 matmul_1、rms_norm、softmax。

对框架开发者来说，这个例子的价值在于：它虽然小，但覆盖了模型里常见的关键模式：

matmul producer
elementwise chain
reduce
broadcast
multi-output use
softmax
second matmul

2. HLO 是什么

HLO 是 High Level Optimizer IR。可以把它理解成：

带 shape / dtype / layout / metadata 的 SSA 张量计算图

SSA 表示每个中间值只定义一次：

%dot = dot(%x, %w1)
%square = multiply(%dot, %dot)
%sum = reduce(%square)
%sqrt = sqrt(...)

HLO 仍然是比较高层的 IR。它关心：

op 类型：dot / reduce / broadcast / exp / sqrt
tensor shape：f32[16,64]
layout：{1,0}
reduce 维度
metadata：来自哪一行 Python，来自哪个 named_call

它暂时不关心：

哪个 MXU 执行 matmul
哪个 cycle 发射指令
何时 vmatpush / vpop
哪个 VLIW bundle 同时发射 DMA 和 vector op

因此，HLO 是模型框架开发者很值得读的层级。很多性能问题在 HLO 层已经能判断：

是否 materialize 了巨大中间矩阵
是否有不必要 broadcast
是否存在 producer 被多个 consumer 使用
是否可以 fusion
layout 是否合理
copy-start / copy-done 是否能 overlap

3. 初始 HLO：从 Python 语义到显式张量图

3.1 Matmul

Python：

h = x @ w1

HLO 中一般表现为：

dot(x, w1)

HLO 会带上 contracting dimensions，例如：

lhs_contracting_dims={1}
rhs_contracting_dims={0}

也就是：

x:  [16,64]
w1: [64,64]
沿 x 的第 1 维和 w1 的第 0 维相乘求和
输出 [16,64]

3.2 RMSNorm-like 部分

Python 直觉：

square = h ** 2
mean = sum(square, axis=-1, keepdims=True) / 64
rms = sqrt(mean + eps)
h_norm = h / rms

HLO 里会显式拆成：

multiply(h, h)
reduce_sum(axis=1)
multiply by 1/64 或 divide by 64
add eps
sqrt
broadcast
divide

注意 broadcast：Python 写 keepdims=True 或隐式 broadcasting 时看起来很自然，但 HLO 必须显式说明哪个 shape 扩展到哪个 shape。

3.3 Softmax

Python：

m = max(h, axis=-1, keepdims=True)
e = exp(h - m)
p = e / sum(e, axis=-1, keepdims=True)

HLO：

reduce_max
broadcast max
subtract
exp
reduce_sum
broadcast sum
divide

HLO 中 reduce 通常会带一个小 computation region，表示规约函数是 add 还是 maximum。

4. HLO optimization passes 做了什么

文章展示了几类重要优化。

4.1 Algebraic simplification

例如：

x / 64

会变成：

x * 0.015625

因为乘法通常比除法更便宜。这类优化在 HLO 层完成，不是 TPU 特有。

4.2 Layout assignment

HLO shape 可能出现：

f32[16,64]{1,0}

解释：

f32[16,64]  dtype + shape
{1,0}       physical layout order

{1,0} 可以粗略理解为最后一维更连续，接近 row-major 的直觉。

4.3 TPU tiling

优化后可能出现：

f32[16,64]{1,0:T(8,128)}

T(8,128) 是 TPU tile annotation。它和 TPU VPU 的向量组织有关：常见向量寄存器结构可以理解为 8 个 sublanes，每个 sublane 128 lanes。

这层含义是：

HLO tensor 不是只作为抽象 ndarray 存放，
backend 已经开始决定怎样把它切成更贴近 TPU 硬件的 tile。

4.4 Fusion

Fusion 是第一篇的主角之一。

没有 fusion：

op A 产生中间 tensor
写出去
op B 读回来
产生另一个中间 tensor
写出去
op C 再读回来

有 fusion：

op A / B / C 进入同一个 fusion kernel
中间值尽量留在寄存器、VMEM 或临时片上状态
只写 fusion 的对外输出

文章里的最终 HLO 大致被压成几个 fusion：

multiply_reduce_fusion:
  matmul_1 + square + reduce_sum
  -> (sum_of_squares, matmul_result)

add_sqrt_fusion:
  sum_of_squares * 1/64 + eps -> sqrt

fusion.5:
  normalized h -> reduce_max

fusion.2:
  exp(h - max) -> reduce_sum

fusion:
  normalize softmax + matmul_2 -> output

5. Multi-output fusion：为什么重要

第一次 matmul 的结果：

h = x @ w1

有两个用途：

用途 1：h -> h ** 2 -> reduce_sum，用于 RMSNorm
用途 2：h -> 后续 normalize / softmax

如果处理不好，编译器可能遇到两难：

要么重复计算 matmul
要么把 matmul result 写到大 buffer，后续再读

multi-output fusion 的做法是：

multiply_reduce_fusion -> (reduce_sum(h*h), h)

它让 matmul 只算一次，同时产出：

1. 每行平方和
2. matmul result 本身

对模型框架开发者来说，这对应一个常见图优化问题：

一个 producer 有多个 consumer。

好的优化器需要在以下选择间做代价判断：

duplicate producer
materialize producer
multi-output fusion
recompute

6. Memory assignment：HBM、VMEM、copy-start、copy-done

文章中 HLO 会出现类似 memory space 的 annotation。粗略理解：

S(0): HBM，片外大内存，通常是默认
S(1): VMEM，片上 SRAM
S(2): sync token 或 backend-specific space

HBM 和 VMEM 的关系：

HBM:
  大，慢，跨 kernel 稳定可见

VMEM:
  小，快，片上，适合 tile 和短生命周期中间值

HLO 中还会出现：

copy-start(w1)
copy-done(w1)

含义：

copy-start:
  发起从 HBM 到 VMEM 的异步拷贝

copy-done:
  等待这个拷贝完成，之后才能使用

这允许 backend 做 overlap：

一边计算当前阶段
一边 DMA 搬下一阶段需要的数据

例如：

先 copy-start(w1)
copy-done(w1) 后调用第一个 matmul fusion

尽早 copy-start(w2)
在 RMSNorm / softmax 期间后台搬 w2
最后 matmul 前 copy-done(w2)

这就是编译器自动做的 prefetch / overlap。

7. Fusion 结束是否一定写回 HBM

这是非常容易误解的点。

更准确的说法：

每个 HLO op 结束不一定写 HBM。
fusion 内部中间值通常不写 HBM。
fusion 的对外输出必须 materialize 到某个 buffer。
这个 buffer 可能是 HBM，也可能是 VMEM，取决于 memory assignment、大小、生命周期和后端能力。

因此：

不是“每个 op 写 HBM”
也不是“每个 fusion 一定写 HBM”
而是“fusion/custom-call 边界通常是 materialization 边界”

在第一篇 toy program 里，一些中间 buffer 可以放到 VMEM。TLP 可以把 VMEM buffer 地址传给后续 fusion，所以并不是所有 fusion output 都落 HBM。

但在第二篇 naive attention 中，scores 形状是：

[heads, seq_len, seq_len]

具体例子：

[8, 2048, 2048] f32
= 8 * 2048 * 2048 * 4 bytes
= 128MB

这个中间矩阵太大，而且跨多个 fusion 被使用：

fusion.5 产生 scores
fusion.2 需要 scores 来做 exp 和 reduce_sum
final fusion 需要 softmax 后的信息继续乘 V

所以它通常会 materialize 到 HBM。第二篇的性能问题本质就是这里。

判断规则：

小的、短生命周期的、编译器能 schedule 在片上的值：
  可能留在 VMEM / register / scratch。

大的、跨 fusion 使用的、生命周期较长的 tensor：
  往往落 HBM。

Flash/Splash Attention 的胜利：
  不是让 128MB scores 写得更快，
  而是让完整 scores tensor 不存在。

8. 从 HLO 到 LLO：层级突然变低

HLO 还像数学图：

dot
reduce
sqrt
exp
fusion

LLO 开始像硬件动作：

vld
vst
vmatpush
vmatmul
vpop
vxpose
vrot.slane
dma.hbm_to_vmem

LLO 是 TPU-specific Low Level Operators。你不需要把每条指令都背下来，但要能识别几个模式。

8.1 TPU 硬件单元粗略图

单元 作用
MXU Matrix Multiply Unit，负责矩阵乘，TPU 主要算力来源
VPU Vector Processing Unit，负责 elementwise、vector add/mul/select 等
XLU 负责 transpose、shuffle、cross-lane movement
Scalar Unit 标量计算、地址、控制流
DMA HBM 与 VMEM 之间异步搬运

8.2 LLO 中识别 matmul

常见模式：

vld        从 VMEM load tile
vmatpush   把 tile push 到 MXU
vmatmul    触发矩阵乘
vpop       从 MXU 取回 accumulator / result

如果看到：

vmatpush
vmatmul
vpop.mrf / vpop.f32

基本可以判断这里在做 MXU matmul。

8.3 LLO 中识别 vector compute

常见：

vadd
vmul
vsub
vsel
vcmp
vpow2
vrsqrt

这些通常是 VPU 上的向量操作。比如 RMSNorm 的 sqrt，后端可能用 vrsqrt 先算 reciprocal sqrt，再组合得到需要的结果。

8.4 LLO 中识别 reduction / transpose

RMSNorm 要做：

sum(h ** 2, axis=-1)

TPU 的 lane/sublane 数据布局和 Python 行列不完全一致，所以 reduction 之前常会出现：

vxpose
vpop.trf

然后是 tree reduction：

vadd
vrot.slane 4
vadd
vrot.slane 2
vadd
vrot.slane 1
vadd

这是并行规约：8 个 sublane 求和，不是串行加 7 次，而是 3 轮合并。

9. VLIW bundle 的特点

VLIW 是 Very Long Instruction Word。

一个 bundle 可以包含多条指令：

bundle:
  dma.hbm_to_vmem
  vld
  vld
  vmov
  scalar address op
  maybe vmatmul

它的核心特点：

1. bundle 内部的指令可以并行发射。
2. 编译器静态决定哪些指令能放进同一个 bundle。
3. 硬件不依赖复杂乱序执行，而是按编译器安排执行。
4. 一个 bundle 可以同时使用多个硬件单元，例如 DMA、MXU、VPU、Scalar。
5. 好的 schedule 会把数据搬运、地址计算、当前 tile 计算、下一 tile 预取交织起来。

所以看到 bundle 时，不要把它理解成：

bundle 中的指令一条条串行执行

更应该理解成：

这是编译器打包好的并行发射包。

文章第一篇说明，toy program 最终被拆成多个 fusion 和一个 TLP；multiply_reduce_fusion 这样的 kernel 会生成几十个 bundle，TLP 负责把多个 fusion 和 DMA 调度起来。

bundle count 有参考意义，但不是唯一性能指标：

bundle 少不一定快
bundle 多不一定慢
HBM traffic、MXU 利用率、DMA overlap、stall、register pressure 都会影响性能

10. TLP：Top Level Program

TLP 是整个 compiled computation 的总控程序。它不是某个 fusion 的内部实现，而是调度多个 fusion/custom-call 和 DMA。

第一篇中可以粗略理解为：

copy w1 HBM -> VMEM
call multiply_reduce_fusion

start copy w2 HBM -> VMEM

call add_sqrt_fusion
call reduce_max fusion
call exp/reduce_sum fusion

wait w2 copy done
call final matmul fusion

final sync

对框架开发者来说：

HLO graph:
  决定有哪些 computation 节点和依赖。

HLO fusion:
  决定哪些 op 合成 kernel。

TLP:
  决定 kernel 和 DMA 的 top-level schedule。

LLO/VLIW:
  决定 kernel 内部如何使用 TPU 硬件单元。

第二篇：When XLA Isn’t Enough

11. 第二篇的问题意识

第一篇告诉我们：

XLA/TPU compiler 能把普通 JAX 程序优化得很深。

第二篇问：

既然 XLA 这么强，为什么还需要 Pallas？

答案：

XLA 能优化你写出来的计算图，
但通常不会自动把 naive attention 改写成 FlashAttention / Splash Attention。

这是“图优化”和“算法改写”的边界。

12. Naive attention 的问题

标准 attention：

scores = Q @ K.T
scores = scores + mask
m = max(scores, axis=-1, keepdims=True)
e = exp(scores - m)
l = sum(e, axis=-1, keepdims=True)
p = e / l
out = p @ V

核心中间矩阵：

scores: [heads, q_len, kv_len]

文章使用的例子中：

heads = 8
seq_len = 2048
head_dim = 128

所以：

scores = [8, 2048, 2048] f32
       = 128MB

XLA 会做真正的优化，例如：

fusion.5:
  Q @ K^T + mask + reduce_max
  -> (max, scores)

fusion.2:
  exp(scores - max) + reduce_sum
  -> sum

fusion:
  normalize + S @ V
  -> output

这里的 XLA 已经很努力：

matmul 和 mask / max 融合
exp 和 sum 融合
normalize 和 final matmul 融合

但问题仍然存在：

完整 scores matrix 仍然被创建出来，并跨 fusion 使用。

所以 attention 的瓶颈不只是“op 没 fusion”，而是：

算法本身 materialize 了巨大的 [H,S,S] 矩阵。

13. 为什么不能简单依赖 XLA

XLA 的优化通常是在等价计算图范围内做：

融合 producer/consumer
消除冗余 broadcast
代数化简
layout assignment
tile scheduling
memory placement

但 FlashAttention / Splash Attention 本质上做了更深的事情：

把 attention 的执行方式改成 streaming over KV blocks。

换句话说：

naive:
  先完整生成 scores
  再 softmax
  再乘 V

Splash:
  每次只生成一个 scores tile
  边扫 KV tile 边维护 online softmax 状态
  从不生成完整 scores

这属于算法级重写。普通 XLA pass 很难从任意 attention graph 自动证明并改写成这种形式，尤其还要考虑 mask、sparsity、GQA/MQA、numerical stability、tile size、memory constraints。

14. 先解决一个直觉误区：KV tile 结果为什么是累加，不是 cat

对一个 query 向量 q，attention 是：

s_j = q @ k_j
p_j = softmax(s)_j
out = sum_j p_j * v_j

如果 KV 被切成两个 tile：

tile 0: k0, k1, k2
tile 1: k3, k4, k5

那么 scores 可以概念上拼接：

scores_tile0 = [s0, s1, s2]
scores_tile1 = [s3, s4, s5]
scores = cat([scores_tile0, scores_tile1])

但是最终输出不是 scores。最终输出是：

out =
  p0 * v0 + p1 * v1 + p2 * v2
  + p3 * v3 + p4 * v4 + p5 * v5

按 tile 写：

out =
  sum_{j in tile0} p_j * v_j
  + sum_{j in tile1} p_j * v_j

所以：

scores 可以在 KV 维度 cat。
output 是沿 KV 维度 weighted sum，因此不同 KV tile 对 output 的贡献要相加。

如果把每个 tile 的 output cat 起来，shape 都错了。对每个 query，attention 最终只输出一个 head_dim 向量。

15. Online softmax：为什么需要 m、l、o

先只看一个 query。它要 attend 到很多 key/value：

key/value 位置：1, 2, 3, ..., N
score：a_1, a_2, a_3, ..., a_N
value：v_1, v_2, v_3, ..., v_N

标准 attention 输出是：

\[\mathrm{Out} = \sum_{t=1}^{N} \frac{e^{a_t}}{\sum_{r=1}^{N} e^{a_r}} v_t\]

为了数值稳定，通常减去最大值：

\[M = \max_{1 \le t \le N} a_t\]

于是：

\[\mathrm{Out} = \frac{ \sum_{t=1}^{N} e^{a_t - M} v_t }{ \sum_{t=1}^{N} e^{a_t - M} }\]

这里可以拆出两个量：

\[O = \sum_{t=1}^{N} e^{a_t - M} v_t\] \[L = \sum_{t=1}^{N} e^{a_t - M}\]

最后：

\[\mathrm{Out} = \frac{O}{L}\]

所以：

L = softmax denominator，分母，也可以理解成总权重
O = softmax @ V 的 numerator，分子，也就是未归一化的加权 V 总和

注意 shape：

L:
  对每个 query 是一个标量。
  对一个 Q block 是 [bq] 或 [bq, 1]。

O:
  对每个 query 是一个 head_dim 向量。
  对一个 Q block 是 [bq, head_dim]。

15.1 分块之后，local_sum 和 local_out 是什么

现在把 KV 分成两个 block：

block A: 位置 1,2,3
block B: 位置 4,5,6

如果已经处理完 block A，我们维护：

\[m_A = \max(a_1,a_2,a_3)\] \[L_A = e^{a_1-m_A} + e^{a_2-m_A} + e^{a_3-m_A}\] \[O_A = e^{a_1-m_A}v_1 + e^{a_2-m_A}v_2 + e^{a_3-m_A}v_3\]

这里：

m_A = 已处理 blocks 的最大 score
L_A = 已处理 blocks 的分母累计
O_A = 已处理 blocks 的加权 V 分子累计

现在来了 block B：

\[m_B = \max(a_4,a_5,a_6)\]

新的全局最大值是：

\[m_{AB} = \max(m_A, m_B)\]

block B 自己的贡献必须用新的最大值 (m_{AB}) 来算：

\[L_B = e^{a_4-m_{AB}} + e^{a_5-m_{AB}} + e^{a_6-m_{AB}}\] \[O_B = e^{a_4-m_{AB}}v_4 + e^{a_5-m_{AB}}v_5 + e^{a_6-m_{AB}}v_6\]

这两个就是代码里的：

local_sum = L_B
local_out = O_B

也就是：

local_sum:
  当前 KV block 对 softmax 分母的贡献。

local_out:
  当前 KV block 对 attention 输出分子的贡献。

对应到 tile 写法：

scores_block = Q_block @ K_block^T
p_block = exp(scores_block - m_new)

local_sum = reduce_sum(p_block, axis=KV_block_dim)
local_out = p_block @ V_block

15.2 为什么旧的 O/L 要乘 correction

问题是：

旧的 L_A / O_A 是用旧最大值 m_A 算的。
现在新的最大值变成了 m_AB。

为了把旧贡献换到新的 max 坐标系，需要缩放：

\[\alpha = e^{m_A - m_{AB}}\]

于是：

\[L_A' = \alpha L_A\] \[O_A' = \alpha O_A\]

这是因为：

\[e^{a_t - m_{AB}} = e^{a_t - m_A} \cdot e^{m_A - m_{AB}}\]

所以旧累计值都要乘同一个缩放因子。

15.3 l_new 和 o_new 是什么

合并旧 blocks 与当前 block：

\[L_{AB} = \alpha L_A + L_B\] \[O_{AB} = \alpha O_A + O_B\]

这就是代码：

l_new = correction * l_prev + local_sum
o_new = correction * o_prev + local_out

变量对照：

correction = alpha
l_prev     = L_A
o_prev     = O_A
local_sum  = L_B
local_out  = O_B
l_new      = L_AB
o_new      = O_AB

最后扫完所有 KV blocks：

\[\mathrm{Out} = \frac{O_{\mathrm{final}}}{L_{\mathrm{final}}}\]

也就是：

output = o_final / l_final

一句话：

Attention = 加权 value 之和 / 权重之和

O = 加权 value 之和，也就是分子
L = 权重之和，也就是分母

online softmax 只是分块处理时，边走边维护 O 和 L。

16. Pallas 是什么

Pallas 是 JAX 的 custom kernel language。官方文档将它描述为 JAX extension，用于为 GPU/TPU 写 custom kernels，同时保留一部分 JAX tracing 和 jax.numpy 风格。Pallas API 仍然是 experimental。

简化理解：

JAX:
  写 tensor program。

Pallas:
  写 tile program。

Pallas 不是手写汇编。你不直接写 VLIW 指令，而是显式表达：

grid 怎么划分
每个 program 处理哪个 tile
输入输出如何 block
scratch 如何跨 iterations 保存
哪些 metadata 放 SMEM

后端仍然负责：

MXU/VPU 映射
DMA scheduling
VMEM layout
VLIW bundle packing

17. Pallas kernel 语法核心

一个典型 Pallas kernel 长这样：

def kernel(q_ref, k_ref, v_ref, o_ref, scratch_ref):
    q = q_ref[...]
    k = k_ref[...]
    v = v_ref[...]

    result = ...

    o_ref[...] = result

这里的参数不是普通 JAX array，而是 Ref。

Ref 可以理解成：

kernel 内部看到的一块可读写内存视图

读取：

x = x_ref[...]

写入：

o_ref[...] = y

调用时用：

pl.pallas_call(
    kernel,
    out_shape=...,
    grid=...,
    in_specs=...,
    out_specs=...,
    scratch_shapes=...,
    compiler_params=...,
)(q, k, v, ...)

核心概念表：

概念 作用
Ref kernel 内部的可读写内存视图
grid kernel program 的多维迭代空间
BlockSpec 每个 grid point 如何映射到 input/output 的 tile
index_map 从 grid indices 返回 array block indices
scratch_shapes 跨 grid iteration 持久存在的临时 buffer
memory_space 指定 VMEM / SMEM 等
compiler_params 给 TPU backend 的编译 hint

18. grid 怎么理解

文章中的 Splash grid 可以抽象成：

grid = (num_heads, num_q_blocks, num_kv_blocks)

一个 grid point：

(h, i, j)

表示：

h: 第几个 attention head
i: 第几个 Q block
j: 第几个 KV block

概念上像三层循环：

for h in range(num_heads):
    for i in range(num_q_blocks):
        initialize m, l, o for this (h, i)

        for j in range(num_kv_blocks):
            scores = Q[h, i] @ K[h, j].T
            update m, l, o

        O[h, i] = o / l

实际 Pallas/TPU backend 可以根据 dimension semantics 和 pipeline 策略重排或并行化某些维度。

19. dimension_semantics 怎么理解

文章提到类似：

dimension 0: heads       -> parallel
dimension 1: q_blocks    -> arbitrary
dimension 2: kv_blocks   -> arbitrary

先不要把它理解成 tensor shape 的维度。它是 grid 的维度。

grid = (h, i, j)

因此：

dimension 0 = h = heads
dimension 1 = i = q_blocks
dimension 2 = j = kv_blocks

19.1 heads 维为什么 parallel

不同 attention head 之间没有依赖：

head 0 不需要 head 1 的 m/l/o
head 1 不需要 head 0 的 m/l/o

所以可以告诉编译器：

heads 维 iteration 可以并行或自由调度。

19.2 q_blocks 维为什么数学上独立

不同 Q block 也有自己的输出和 scratch：

Q block 0 有自己的 m/l/o
Q block 1 有自己的 m/l/o

数学上它们是独立的。不过文章示例中可能仍把它标成 arbitrary，这是保守选择，表示不要让编译器对这一维做过强假设。

19.3 kv_blocks 维为什么有依赖

对固定 (h, i)，要沿 j 扫过 KV blocks：

j = 0 -> 得到 m0, l0, o0
j = 1 -> 需要读 m0, l0, o0，更新成 m1, l1, o1
j = 2 -> 需要读 m1, l1, o1，更新成 m2, l2, o2

所以 KV block 维不能随便并行或乱序，因为 online softmax 的状态沿这一维传递。

一句话：

parallel across heads
independent across Q blocks
accumulate across KV blocks

20. BlockSpec 怎么理解

BlockSpec 定义：

某个 grid point 应该看到数组的哪一个 tile。

例如 Q：

pl.BlockSpec(
    block_shape=(None, bq, head_dim),
    index_map=lambda h, i, j, *_: (h, i, 0),
)

假设 Q shape 是：

Q: [num_heads, q_len, head_dim]

index_map 返回：

(h, i, 0)

分别对应：

axis 0: head 轴       -> 第 h 个 block
axis 1: sequence 轴   -> 第 i 个 Q block
axis 2: head_dim 轴   -> 第 0 个 feature block

结合：

block_shape=(None, bq, head_dim)

可以近似理解为：

Q[h, i*bq:(i+1)*bq, 0:head_dim]

注意：

None 不是“读完整维度”。
None 表示该维度取 size-1 的 slice，并在 kernel 内部 squeeze 掉。

所以 kernel 内部看到的 Q tile 通常是：

[bq, head_dim]

而不是：

[1, bq, head_dim]

20.1 为什么 lambda h, i, j: (h, j, 0) 中的 0 表示读完整 head_dim

它不是因为 0 有“完整”的特殊含义。

对于 K：

K: [num_kv_heads, kv_len, head_dim]

BlockSpec：

pl.BlockSpec(
    block_shape=(None, bkv, head_dim),
    index_map=lambda h, i, j: (h, j, 0),
)

真实 slice 近似是：

K[h, j*bkv:(j+1)*bkv, 0*head_dim:1*head_dim]

也就是：

K[h, j*bkv:(j+1)*bkv, 0:head_dim]

所以：

0 只是 head_dim 这一轴的 block index。
因为 block_shape 的最后一维恰好等于完整 head_dim，
所以从第 0 个 block 开始读，就读到了完整 feature 维。

如果 block_shape 写成：

block_shape=(None, bkv, 32)

而 head_dim=128，那么：

0 -> 读 0:32
1 -> 读 32:64
2 -> 读 64:96
3 -> 读 96:128

此时 0 就不代表完整 head_dim 了。

21. K/V 的 sparse indirection

Dense attention 里，K/V 的 index map 可以很简单：

lambda h, i, j: (h, j, 0)

意思是：

当前 program 是 (h, i, j)
就读取第 h 个 head、第 j 个 KV block

Sparse attention 里，某些 KV block 对当前 Q block 完全被 mask，没必要加载和计算。Splash 使用 data_next_ref 做间接寻址：

def k_index_map(h, i, j, data_next_ref, block_mask_ref, mask_next_ref):
    next_j, *_ = _next_nonzero(
        h, i, j,
        data_next_ref,
        block_mask_ref,
        mask_next_ref,
    )
    return (h // q_heads_per_kv_head, next_j, 0)

直觉：

本来要读 KV block j
但如果 j 是 fully masked block
就通过 data_next_ref 跳到下一个有效 block next_j

例如：

j:       0  1  2  3  4  5
valid:   1  0  0  1  0  1

那么：

data_next[h,i,1] 可能指向 3
data_next[h,i,2] 也可能指向 3

这样 kernel 在寻址阶段就跳过无效块：

不是 load 之后发现不用，
而是根本不 load 那些 fully masked KV blocks。

h // q_heads_per_kv_head 用于 GQA/MQA：

Q heads 可能比 KV heads 多。
多个 Q heads 共享同一个 KV head。

例如：

Q heads = 8
KV heads = 2
q_heads_per_kv_head = 4

则：

Q head 0,1,2,3 -> KV head 0
Q head 4,5,6,7 -> KV head 1

22. Pallas 暴露的 TPU memory hierarchy

文章中的核心内存概念：

HBM:
  片外大内存。
  Q/K/V/O 原始大数组通常在这里。

VMEM:
  片上 SRAM。
  默认 BlockSpec refs 指向 VMEM tile。

SMEM:
  scalar memory。
  适合小索引、mask metadata、next pointer、control-flow decision。

Scratch:
  kernel iterations 之间持久存在的临时 buffer。

22.1 VMEM：BlockSpec refs 默认在这里

例如：

pl.BlockSpec((bq, head_dim), index_map)

外面的 Q/K/V 大数组在 HBM，进入 kernel 时当前 tile 被搬到 VMEM。kernel 内部看到的是：

Q tile: [bq, head_dim]
K tile: [bkv, head_dim]
V tile: [bkv, head_dim]

22.2 SMEM：放小的控制信息

例如：

pl.BlockSpec(
    (num_heads,),
    lambda *_: (0,),
    memory_space=pltpu.SMEM,
)

适合：

block mask metadata
data_next pointer
mask_next pointer
小整数索引
控制流判断

这些不是大矩阵数据，而是告诉 kernel：

当前 block 是否有效
下一个有效 block 是谁
怎样跳过 masked region

22.3 Scratch：online softmax 的记忆

文章中的：

m_scratch_ref
l_scratch_ref
o_scratch_ref

分别保存：

m: running max
l: running denominator
o: running unnormalized output numerator

对固定 (h, i)，沿 j 扫 KV blocks：

j=0:
  初始化并写 m/l/o scratch

j=1:
  读旧 m/l/o scratch
  更新成新的 m/l/o

j=2:
  继续

最后:
  output = o / l
  写最终 O

关键：

这些 scratch 在 VMEM 里跨 grid iterations 保留，
不需要每个 KV tile 都把 m/l/o 写回 HBM 再读回来。

这就是 online softmax 能高效实现的关键之一。

23. Splash Attention 的执行逻辑

对一个 Q block 和一个 KV block：

Q_tile: [bq, head_dim]
K_tile: [bkv, head_dim]
V_tile: [bkv, head_dim]

当前 tile scores：

scores_tile = Q_tile @ K_tile.T

shape：

[bq, bkv]

这个 scores tile 只在 VMEM 中短暂存在。然后：

更新 running max m
更新 running denominator l
更新 running output numerator o

再处理下一个 KV tile。

最终：

O_tile = o / l

写回 HBM。

对比 naive：

naive:
  scores_full = Q @ K.T
  scores_full shape = [H, S, S]
  scores_full 会跨 fusion materialize

Splash:
  scores_tile = Q_tile @ K_tile.T
  scores_tile shape = [bq, bkv]
  scores_tile 不跨 kernel 边界

这就是 HBM traffic 大幅下降的原因。

24. Pallas kernel 到 HLO：为什么是 custom-call

Pallas kernel 在 HLO 中通常表现为：

custom-call(...)
custom_call_target="tpu_custom_call"

这意味着：

XLA HLO graph 不再展开 Pallas kernel 内部细节。

XLA 看到的是：

这里有一个 custom-call
输入是 Q/K/V/mask metadata
输出是 O

Pallas kernel body 作为 MLIR payload 交给 Mosaic。之后 Mosaic 和 TPU backend 继续降到：

LLO -> VLIW bundles

所以 Pallas 的定位不是：

绕开整个 TPU compiler

而是：

绕开 XLA 高层图优化对 kernel 内部算法的表达限制，
但仍然使用 Mosaic / TPU backend 做底层 lowering 和 scheduling。

25. 两篇文章的最终对比

维度 普通 JAX / XLA Pallas / Splash
你写的东西 高层 tensor program tiled kernel program
XLA 是否看见内部 看见完整 HLO graph HLO 只看见 custom-call
优化主要来自 fusion、layout、tiling、memory assignment、VLIW scheduling 你显式表达算法级 tiling / streaming，后端继续调度
attention scores 完整 [H,S,S] tensor 存在 只有 [bq,bkv] tile 短暂存在
softmax 对完整 scores 分阶段处理 online softmax，维护 m/l/o
中间矩阵 HBM traffic 大 小
程序员负担 低，主要写 JAX 高，需要理解 grid、BlockSpec、scratch、memory space
适用场景 常规模型算子、可由 XLA 表达的图优化 编译器不会自动发现的算法级优化

26. 给模型框架开发者的复习清单

26.1 HLO 层重点看什么

读 HLO 时优先看：

1. 是否出现巨大中间 tensor
   例如 attention scores [H,S,S]

2. producer/consumer 是否 fusion
   elementwise + reduction + broadcast 是否合并

3. common producer 如何处理
   multi-output fusion / duplicate / materialize / recompute

4. layout 是否合理
   是否出现 TPU tile annotation

5. memory assignment
   大 tensor 是否 HBM
   小中间结果是否 VMEM

6. copy-start/copy-done
   是否有异步拷贝和 compute overlap

26.2 LLO 层只需先会识别模式

vmatpush / vmatmul / vpop
  -> MXU matmul

vadd / vmul / vsel / vcmp / vpow2 / vrsqrt
  -> VPU vector compute

vxpose / vpop.trf / vrot.slane
  -> transpose / shuffle / reduction

dma.hbm_to_vmem / dma.done.wait
  -> HBM 与 VMEM 之间数据搬运

bundle { ... }
  -> VLIW 静态并行发射包

26.3 Pallas 层重点看什么

1. grid
   每个 program id 代表什么？
   哪些维度 independent，哪些维度 carry state？

2. BlockSpec
   每个 grid point 看到数组的哪个 tile？

3. index_map
   是否存在 sparse indirection？
   是否有 GQA/MQA head mapping？

4. scratch
   哪些状态跨 grid iteration 保留？
   是否避免 HBM round-trip？

5. memory_space
   哪些 metadata 放 SMEM？
   哪些 tile 放 VMEM？

6. custom-call
   HLO 是否只看见 opaque call？
   kernel body 是否进入 Mosaic？

27. 最后用一句话复习

第一篇：

JAX 写自然的 tensor program，XLA 可以把它优化成带 fusion、VMEM placement、DMA overlap 和 VLIW bundle 的 TPU 程序。

第二篇：

当性能瓶颈来自算法级 materialization，例如完整 attention matrix，XLA 的 fusion 仍然不够；Pallas 让我们直接表达 streaming tiled algorithm，用 online softmax 避免完整 scores 落 HBM。

最核心判断：

如果问题是“这些 op 能不能合并得更好”，先看 XLA/HLO fusion。
如果问题是“这个巨大中间 tensor 是否本来就不该存在”，考虑算法改写和 Pallas。

参考资料

]]>MoE 优化的探索2026-05-10T00:00:00+00:002026-05-10T00:00:00+00:00https://wqh011128.github.io/blog/2026/05/10/MoE-OptimizationMoE 优化的探索：从 MiniMax-01、Comet、FlashMoE 到 DeepSeek-V4

Link:

[toc]

Main Idea

MoE 的核心收益来自 sparse computation：每个 token 只激活少数 experts，因此模型总参数可以很大，但单 token 计算量仍然可控。

问题是，MoE 在分布式训练和推理里会引入很重的通信。一个 token 被 router 分到某个 expert 后，那个 expert 可能在另一张 GPU 上，于是系统必须先把 token 发过去，再把 expert 的输出发回来。这个过程通常叫：

Dispatch -> Expert Compute -> Combine

当模型规模变大时，瓶颈不再只是 expert 的矩阵乘，而是 通信和计算之间的空转。如果所有 GPU 都先等 Dispatch All-to-All 完成，再开始 expert GEMM，算完之后又一起等 Combine All-to-All，那么 GPU 会在通信阶段闲着，网络也会在计算阶段闲着。

MoE 优化的主线就是：

不要让通信和计算串行排队，而是把 token、expert、tile 或 kernel 切成更细粒度，让“正在算的一部分”和“正在传的一部分”重叠起来。

1. MoE 基本流程

以常见的 Transformer MoE FFN 为例，一个 MoE 层可以分成四类操作。

阶段 做什么 主要瓶颈
Router 为每个 token 选择 top-k experts 负载均衡、路由开销
Dispatch 把 token 发到 expert 所在设备 all-to-all 通信
Expert FFN 每个 expert 对收到的 token 做 MLP GEMM 计算
Combine 把 expert 输出发回原位置并加权合并 all-to-all 通信

Expert FFN 本身通常又包含两个大 Linear：

Linear-1: X_e -> hidden expansion
Activation: SiLU / SwiGLU
Linear-2: hidden expansion -> output

如果写成 SwiGLU 风格：

\[\begin{aligned} g_e &= X_e W_{e,\mathrm{gate}} \\ u_e &= X_e W_{e,\mathrm{up}} \\ z_e &= \mathrm{SiLU}(g_e) \odot u_e \\ Y_e &= z_e W_{e,\mathrm{down}} \end{aligned}\]

其中 $X_e$ 是分配给 expert $e$ 的 token 表示。Linear-1 可以理解为 $W_{e,\mathrm{gate}}$ 和 $W_{e,\mathrm{up}}$ 这一侧的 GEMM，Linear-2 是 $W_{e,\mathrm{down}}$ 的 GEMM。

这里的 GEMM 是 General Matrix Multiplication，也就是通用矩阵乘。GPU 对 GEMM 极度优化，所以 expert compute 的核心就是如何让这些 GEMM 一直吃满算力。

2. MiniMax-01：token group / process group 级别的 overlap

在MiniMax-01的总结中解读了EP，EP overlap，以及MiniMax的MoE创新。其优化重点不是把 expert 内部拆成 Linear-1 / Linear-2 的流水线，而是围绕 Expert Parallelism (EP)、Expert Tensor Parallelism (ETP) 和 Expert Data Parallelism (EDP) 做更合理的通信计算重叠。

2.1 Token-grouping-based overlap

MiniMax-01 先把 tokens 切成多个 group。每个 group 都要经历：

a2a-dispatch -> expert compute -> a2a-combine

如果完全串行，时间线是：

group 0: dispatch -> compute -> combine
group 1: dispatch -> compute -> combine
group 2: dispatch -> compute -> combine

MiniMax-01 希望改成：

group 0: dispatch -> compute -> combine
group 1:          dispatch -> compute -> combine
group 2:                   dispatch -> compute -> combine

这样，某个 group 在做 expert compute 时，另一个 group 可以做 dispatch 或 combine。

这是一种 token group 粒度 的 overlap。它没有改变 expert FFN 的数学结构，也没有把单个 expert 的 GEMM 拆成 wave（像DeepSeek v4那样）。它只是把 token batch 切小，让通信和计算不要在整层级别完全串行。

2.2 为什么还要 ETP / EDP

MiniMax-01 进一步指出，仅靠 EP 不一定够。当 expert 参数太大时，可以用 Expert Tensor Parallelism (ETP) 把单个 expert 的参数也切到多个设备上。

这时 MoE 层的流程会变成：

a2a-dispatch -> allgather -> expert compute -> reduce-scatter -> a2a-combine

这里的 allgather 和 reduce-scatter 来自 tensor parallelism：为了让多个设备共同计算一个 expert，需要先收集输入或中间张量，再把结果规约分散回去。

MiniMax-01 还引入 Expert Data Parallelism (EDP)，本质上是在 expert 维度做复制，缓解某些 expert 负载过高的问题。

2.3 MiniMax-01 的定位

MiniMax-01 更像是训练系统级别的 MoE 并行优化：

  • 重点是 EP / ETP / EDP 的组合。
  • overlap 粒度主要是 token group 和 process group。
  • 它处理的是大规模训练时不同并行策略之间的通信压力。
  • 它不是 Comet / FlashMoE / DeepSeek 那种更细的 expert GEMM pipeline 或 kernel-level 调度。

3. Comet：fine-grained computation-communication overlapping

这一节开始单独精读 Comet: Fine-grained Computation-communication Overlapping for Mixture-of-Experts。这篇文章可以看成是在回答一个很具体的问题：

传统 EP overlap 已经把 token batch 切成 chunk 了，为什么 MoE 层里还是有明显 GPU idle time？如果继续往下切，应该切什么，怎么切，谁先算，谁先传？

Comet 的答案不是简单地“再把 chunk 切小一点”。它真正做的是：找到 MoE layer 中通信算子和计算算子之间共享的 buffer，也就是 shared tensor，然后根据 consumer 的依赖关系决定沿哪个维度切分，并重新安排 GroupGEMM 的 tile 执行顺序。

3.1 传统 EP overlap 为什么还是 coarse-grained

Comet 在 Introduction 里先分析了传统做法。一个分布式 MoE 层通常可以抽象成：

Receive / Dispatch
-> Expert computation
-> Send / Combine

如果不 overlap，就是先收完所有 token，再做 expert GEMM，最后再发回结果。传统 EP overlap 会把 expert computation kernel 切成几个 chunk，让某个 chunk 的计算和另一个 chunk 的通信同时发生。

Comet Figure 1: MoE execution analysis

Figure 1(b) 想表达的是：把输入拆成 chunk 后，确实可以让一部分通信和一部分计算重叠，但这种 overlap 仍然是 coarse-grained。原因有三个。

第一，chunk 仍然必须作为一个整体 ready。即使 chunk 已经比完整 batch 小，一个 chunk 内部仍然可能有很多 token；只要这个 chunk 需要的 token 还没全部到齐，expert computation 就不能启动。

第二，chunk 变小会伤害 GEMM 效率。论文里提到，原本完整 expert computation 的时间是 $t$，切成两个 chunk 后可能变成 $t_1+t_2>t$。这不是数学计算量变多了，而是小 GEMM 更难吃满 tensor core，还会带来更多调度和访存开销。

第三，MoE 是动态的。这里的动态不是模型结构变了，而是 router 每次会把不同 token 分给不同 experts。某个 step 里 Expert0 可能收到很多 token，Expert1 很少；下个 step 又可能反过来。于是每个 expert 的输入形状、通信量、计算量都在运行时变化，导致“通信 chunk”和“计算 chunk”的时间很难稳定对齐。

所以传统 EP overlap 的核心问题是：它把通信和计算装进不同 kernel / stream 里，让它们粗粒度并行，但对 GPU thread blocks、GEMM tile 顺序、remote I/O 这些底层资源缺少精细控制。

3.2 MoE Structure：论文里的 forward 过程

Comet 的 Figure 2 给了一个普通 MoE layer 的执行例子：两个 GPU，总共四个 experts，GPU0 放 Expert0/Expert1，GPU1 放 Expert2/Expert3。每个 token 被 router 分到 $k$ 个 experts。图里 Token A 被路由到 Expert0、Expert1、Expert3。

Comet Figure 2: MoE layer across two GPUs

论文中几个重要符号可以这样读：

符号 含义
$E$ expert 总数
$k$ 每个 token 被路由到的 expert 数量，也就是 top-k
$TP$ tensor parallel size
$EP$ expert parallel size
$TP \times EP$ 总并行 world size
$M$ GEMM 中的 row 维度，在 MoE 里通常对应 token / token-expert rows
$K$ GEMM reduction 维度，例如输入 hidden 或 FFN intermediate 维
$N$ GEMM output column 维度，例如输出 hidden 或 FFN intermediate 维
$T_M$ GEMM tile 在 $M$ 维的大小
$T_N$ GEMM tile 在 $N$ 维的大小

MoE 的每个 expert FFN 有两层 GEMM：

layer0: first expert GEMM, usually up/gate projection
activation: SiLU / SwiGLU
layer1: second expert GEMM, usually down projection

论文把 MoE 的执行过程分成两类 pipeline。

communication-computation pipeline: 对应 MoE layer0。

Dispatch communication -> layer0 GroupGEMM

这里通信是 producer，layer0 GroupGEMM 是 consumer。shared tensor 是 dispatch 之后、即将被 layer0 GEMM 消费的 expert input buffer。

computation-communication pipeline: 对应 MoE layer1。

layer1 GroupGEMM -> top-k routed reduction / combine communication

这里 layer1 GroupGEMM 是 producer，reduction/combine 是 consumer。shared tensor 是 layer1 GEMM 产生、即将被合并和通信的 expert output buffer。

3.3 TP 是什么，为什么 MoE 会用

Tensor Parallelism (TP) 是把同一个线性层的权重切到多个 GPU 上。它和 EP 的区别是：

EP: 不同 experts 放在不同 GPU 上，每个 expert 权重通常是完整的。
TP: 同一个 expert / linear 的权重沿 hidden dimension 切开，多张 GPU 一起算一个矩阵乘。

比如一个线性层：

\[Y = XW,\quad W\in\mathbb{R}^{K\times N}\]

如果按 $N$ 维做 column parallel，GPU0 负责 $W[:,0:N/2]$，GPU1 负责 $W[:,N/2:N]$。如果按 $K$ 维做 row parallel，多个 GPU 分别计算部分乘积，再做 reduce。

MoE 会用 TP，是因为单个 expert 的 FFN 也可能很大。只用 EP 时，每个 expert 权重完整放在某张 GPU 上；如果 expert 太大，或者希望提升单 expert GEMM 的吞吐，就需要 TP 继续切 expert 内部权重。实际大模型里常常是 EP + TP 混合并行。

3.4 Granularity mismatch：token-level communication vs tile-level computation

Comet 的第一个关键观察是 granularity mismatch between computation and communication。

在 MoE 里，通信的基本单位通常是 token。Router 决定某个 token 要去哪个 expert，于是系统把这个 token 发到 expert 所在 GPU。

但高性能 GEMM 的基本单位不是单 token，而是 tile。论文里 Figure 2 的紫色块就是一个 computation tile，例如 $128\times128$。这意味着一个 expert 的某个 GEMM tile 可能需要 128 个 token rows，而这些 token 由 router 决定，可能随机分布在多个 GPU 上。

这就产生了依赖：

一个 GEMM tile 需要的 token rows 没有全部 ready
-> 这个 tile 不能开始计算

所以问题不是“tile 大小不一样导致完成时间不一样”，而是：

每个 tile 依赖的 token 来源不同，ready 时间不同；coarse-grained dispatch 会让 tile 等整个 chunk 或整个 expert input buffer。

Comet 因此提出 fine-grained communication：每个 computation tile 通过 Unified Virtual Address (UVA) 直接读/写它需要的数据。

UVA 的实际作用是提供统一虚拟地址空间。在支持 GPU peer access 的情况下，GPU kernel 可以拿到远端 GPU buffer 的地址，并发起细粒度 remote load/store。它不是让 tile “自己有智能”，而是让 kernel 里的 communication thread blocks 可以根据路由 metadata，把某个 tile 需要的 remote token rows 拉到本地，或者把某个输出 tile 写回目标位置。

但 fine-grained remote I/O 很慢。如果把远程读写塞进 GEMM compute thread block，会破坏 tensor core pipeline。Comet 后面才需要 thread block specialization：通信 block 专门做 remote I/O，计算 block 保持高效 GEMM。

3.5 Design overview：shared tensor 是桥

Comet 的 Figure 3 是整篇文章的设计总览。

Comet Figure 3: Design overview

Comet 有两个核心设计：

机制 解决什么
Shared tensor based dependency resolving 分析 producer/consumer 之间的真实数据依赖，决定 shared tensor 沿哪个维度切，并重排 tile 顺序
Adaptive workload assignment 在 fused kernel 内动态分配 thread blocks 给通信和计算，减少 pipeline bubble

这里的 shared tensor 可以简单理解成：producer 和 consumer 共用的那块中间 buffer。

对 layer0：

producer: dispatch communication
shared tensor: expert input X_e
consumer: layer0 GroupGEMM

对 layer1：

producer: layer1 GroupGEMM
shared tensor: expert output Y_e
consumer: top-k routed reduction + combine communication

shared tensor 重要，是因为 overlap 只有在 producer 和 consumer 能处理 shared tensor 的不同独立部分时才成立。如果 consumer 必须等完整 tensor，overlap 就退化成普通串行。

3.6 3.1.1：How to decompose the shared tensor

Figure 4 把 layer0 和 layer1 都建模成 producer-consumer 关系。

Comet Figure 4: Producer-consumer modeling

Comet 的原则是：

沿 consumer 视角下相互独立的维度切 shared tensor。

3.6.1 Layer0 为什么沿 $M$ 切

Layer0 的 shared tensor 是 layer0 GEMM 的输入矩阵：

\[X_e \in \mathbb{R}^{M_e\times K}\]

其中 $M_e$ 是 expert $e$ 收到的 token rows 数量，$K$ 是 token embedding / hidden dimension。

Layer0 的 consumer 是 GEMM：

\[H_e = X_e W_{e,0}\]

对 GEMM 来说，不同 token rows 之间相互独立。也就是说，先算 $X_e[M_0,:]$ 和后算 $X_e[M_1,:]$ 不会改变结果。因此 layer0 可以沿 $M$ 维切：

X_e[M0, :] -> layer0 GroupGEMM
X_e[M1, :] -> layer0 GroupGEMM
X_e[M2, :] -> layer0 GroupGEMM

但不能沿 $K$ 维随便切，因为 GEMM 对 $K$ 维做 reduction。算一个输出元素需要完整的 $K$ 维乘加：

\[H_{e,i,n}=\sum_{k}X_{e,i,k}W_{e,0,k,n}\]

如果切 $K$，不同分块之间还要额外做 partial sum reduction，consumer 不能直接独立消费。

3.6.2 Layer1 为什么沿 $N$ 切

Layer1 的 shared tensor 是 layer1 GEMM 的输出：

\[Y_e \in \mathbb{R}^{M_e\times N}\]

Layer1 后面的 consumer 不是普通逐元素操作，而是 top-k routed reduction + combine。注意，这里的 top-k reduction 不是重新选择 top-k；top-k 在 router 阶段已经完成了。这里的意思是：对同一个原始 token 的多个 expert 输出按 router weight 做加权合并。

例如 top-2：

\[O_t = w_{t,e_1}Y_{t,e_1}+w_{t,e_2}Y_{t,e_2}\]

如果沿 $M$ 切，可能把同一个 token 的两个 expert 输出拆到不同块：

M tile 0: token A from Expert0
M tile 1: token A from Expert3

这时 consumer 处理 M tile 0 时拿不到 token A 的完整 top-k routed outputs，因此 $M$ 维存在 interdependency。

但 $N$ 维是 output feature column。不同 feature columns 的 weighted reduction 相互独立：

\[O_{t,n}=w_{t,e_1}Y_{t,e_1,n}+w_{t,e_2}Y_{t,e_2,n}\]

所以 layer1 可以沿 $N$ 切：

Y[:, N0] -> reduction + combine
Y[:, N1] -> reduction + combine
Y[:, N2] -> reduction + combine

这就是论文中“layer0 沿 $M$ 分解，layer1 沿 $N$ 分解”的根本原因。

3.7 3.1.2：How to reschedule the decomposed shared tensor

只知道沿哪个维度切还不够。Comet 还要决定切完之后怎么排执行顺序。论文给了两个原则：

  1. sub-tensors 要尽量对齐原始 GEMM tile granularity，否则 GEMM 效率会下降。
  2. 优先执行 producer 已经产出、consumer 可以立即使用的部分，让 consumer 尽早启动。

3.7.1 Layer0：按 $M$ 切后，先算 local-token tiles

Layer0 是：

Dispatch -> layer0 GroupGEMM

Figure 5 画的是 Rank0 上有三个 experts，每个 expert 都需要 local data 和 remote data。

Comet Figure 5: Decompose and reschedule layer0 shared tensor

Comet 会先按 source rank 对 token 排序。直觉上：

local tokens | remote rank 1 tokens | remote rank 2 tokens | ...

然后 GroupGEMM 的 tile compute sequence 会优先从 local tokens 所在 tile 开始。这样本地 tile 可以马上计算，同时远程 token 还在通过 communication blocks 传输。

时间线可以理解成：

t0: compute local-token tiles
    communicate remote-rank-1 token rows

t1: compute remote-rank-1 tiles
    communicate remote-rank-2 token rows

t2: compute remote-rank-2 tiles

这里的 tile 通常类似 $T_M\times T_N$，例如论文前文举的 $128\times128$。但要注意，这个 tile 是 某个 expert GEMM 内部的 tile，不是把不同 experts 的 token 混在一起乘同一个权重。

3.7.2 关键疑问：不同 experts 权重不同，$128\times128$ tile 怎么算

这是理解 GroupGEMM 的关键。

GroupGEMM 不是把所有 experts 的 token 拼成一个大矩阵，然后乘同一个权重。它是一组独立 GEMM 的调度：

\[H_e = X_e W_{e,0},\quad e\in\mathcal{E}_{\mathrm{local}}\]

也就是说，Rank0 上如果有 Expert0、Expert1、Expert2，那么 GroupGEMM 实际上在执行：

Expert0: X_0 @ W_0
Expert1: X_1 @ W_1
Expert2: X_2 @ W_2

每个 computation tile 都带着 expert id。属于 Expert0 的 tile 只会用 $W_0$，属于 Expert1 的 tile 只会用 $W_1$。高性能 grouped GEMM kernel 会通过 metadata / pointer arrays / offsets 找到每个 expert 对应的 input pointer、weight pointer 和 output pointer。

所以 Figure 5 中的 $M$ 切分不是：

把不同 expert 的 128 行混成一个 tile，用同一个 W 去乘

而是：

在每个 expert 自己的 X_e[M_e, K] 里按 M tile 切；
GroupGEMM 只是把多个 expert 的 tile 放到同一个 kernel 里统一调度。

如果某个 expert 收到的 token 不足一个完整 $T_M$，实现上可以用 partial tile、padding 或 grouped GEMM 的 ragged shape metadata 处理。数学上仍然是每个 expert 使用自己的权重。

这也解释了为什么 layer0 按 $M$ 切完不需要“还原成原始 token 顺序”再进入 layer1。Layer0 输出仍然按 expert 分组保存：

\[H_e[M_e, K']\]

中间 activation 是逐元素的，不需要跨 expert 或跨 token 重新排列。Layer1 继续对同一个 expert 的 $H_e$ 做：

\[Y_e = H_e W_{e,1}\]

真正需要恢复到原始 token 顺序，是 layer1 结束后的 routed reduction / combine 阶段。

3.7.3 Layer1：按 $N$ 切后，column-wise 执行 GroupGEMM

Layer1 是：

layer1 GroupGEMM -> reduction + combine communication

Figure 6 说明了 Comet 如何重排 layer1 的 GroupGEMM。

Comet Figure 6: Rescheduled compute sequence for layer1

如果不重排，GroupGEMM 可能按 expert 顺序执行：

Expert0: N0 -> N1 -> N2 -> N3
Expert1: N0 -> N1 -> N2 -> N3
Expert2: N0 -> N1 -> N2 -> N3

这样 consumer 很难提前开始，因为它想处理某个 column block 时，需要相关 experts 的同一段 columns 都已经产生。

Comet 改成 column-wise：

N0 group:
  Expert0:N0 -> Expert1:N0 -> Expert2:N0

N1 group:
  Expert0:N1 -> Expert1:N1 -> Expert2:N1

N2 group:
  Expert0:N2 -> Expert1:N2 -> Expert2:N2

注意，这里的 N0 不是“对 128 列做 top-k selection”。Top-k selection 已经在 router 完成。这里做的是：对已经确定的 top-k experts，在 N0 这一段 output features 上做 weighted reduction 和 combine。

如果 $N_0$ 表示 columns $0:T_N$，那么 consumer 可以先做：

\[O_{t,N_0}=\sum_{e\in\mathrm{TopK}(t)}w_{t,e}Y_{t,e,N_0}\]

同时 layer1 GroupGEMM 继续计算 $Y[:,N_1]$。这样就形成：

compute Y[:, N0]
-> reduce/combine Y[:, N0]

while

compute Y[:, N1]

所以 Figure 6 的核心不是“列维度上重新选 top-k”，而是：按 column block 提前产出可被 consumer 完整处理的一段 output features。

3.8 3.2：Adaptive Workload Assignment

经过 3.1 的 dependency resolving 后，Comet 已经知道哪些数据可以先算、哪些数据可以先传。但还有一个问题：fine-grained remote I/O 很慢，GEMM 又很吃 tensor core。谁来做通信，谁来做计算，分配多少 GPU 资源，不能拍脑袋。

Figure 7 展示的是 Comet 在 Hopper 上的 fused kernel 设计。

Comet Figure 7: Thread block specialized kernel

3.8.1 Thread block specialization

最直接的融合方式叫 vertical fusion：每个 thread block 既做 GEMM，也在 prologue / epilogue 里做通信 I/O。

问题是 remote I/O 延迟远高于本地显存访问。如果把 remote read/write 插进 GEMM thread block，可能会阻塞后续 tensor core 计算，尤其 Hopper 上 GEMM 通常利用 TMA 建立异步 compute pipeline，长延迟 remote I/O 会破坏这个 pipeline。

Comet 因此把 thread blocks 隔离成两类：

compute thread blocks: 负责 GEMM，尽量复用默认 CUTLASS GEMM 实现
communication thread blocks: 负责 remote I/O、top-k routed reduction、local/remote writeback

这样做的代价是会多一些 global memory 读写，但收益是通信不会污染 GEMM 的关键路径，而且系统可以精确控制多少 blocks 做通信、多少 blocks 做计算。

3.8.2 Adaptive thread block assignment

论文 3.2.2 解决的是：通信 block 和计算 block 到底分多少？

假设一个 fused kernel 总共有 $N_{\mathrm{TB}}$ 个 thread blocks，其中：

N_{\mathrm{comp}}: compute blocks
N_{\mathrm{comm}}: communication blocks
N_{\mathrm{TB}} = N_{\mathrm{comp}} + N_{\mathrm{comm}}

如果 $N_{\mathrm{comm}}$ 太少，远程 I/O 跟不上，GEMM 算完后会等通信。 如果 $N_{\mathrm{comm}}$ 太多，GEMM blocks 变少，计算吞吐下降。 最佳分界点和输入 token length、TP/EP 配置、expert shape、硬件带宽都有关。

Figure 8 说明不同配置下最优 $N_{\mathrm{comm}}$ 不同。

Comet Figure 8: Adaptive thread block assignment

论文给的例子是：当输入 token length 从 4096 变到 16384 时，最优通信 block 数会变化；当 TP 从 8 调到 4 时，最优分配点也会明显变化。

所以 Comet 的做法不是运行时在线搜索，而是：

1. 预编译多个 kernel，每个 kernel 使用不同的 compute/communication block division point。
2. 部署前 profile 不同模型配置和输入形状，记录最优配置 metadata。
3. 运行时根据 metadata 选择合适 kernel。

这就是 adaptive workload assignment。它的目标不是改变数学计算，而是让 fine-grained pipeline 的通信段和计算段时间尽量对齐，减少 pipeline bubbles。

3.9 把我们讨论过的几个易错点放在一起

Tile 不是 token chunk。 token chunk 是 EP overlap 的 coarse 粒度；tile 是 GEMM kernel 的计算分块，例如 $128\times128$。一个 tile 通常包含多个 token rows 和一段 output columns。

Tile 不是 expert 的最大容量。 一个 expert 实际收到多少 token 由 router 决定，记作 $M_e$。GEMM tile 是在 $M_e\times K$ 或 $M_e\times N$ 这个矩阵内部继续切出来的计算块。

Layer0 的 $M$ 切分发生在 layer0 计算前。 它切的是 dispatch 后的 expert input $X_e$，目的是 token rows 到一块、算一块。

Layer1 的 $N$ 切分发生在 layer1 计算过程中。 它不是先完整算出 $Y_e$ 再切，而是调整 GroupGEMM 顺序，直接先产出 $Y[:,N_0]$，让 reduction/combine 提前消费。

Layer0 到 layer1 中间不需要还原成原始 token 顺序。 layer0 输出仍然按 expert 分组，activation 和 layer1 都可以在 expert-local layout 里继续做。只有最终 combine 时才需要根据 routing metadata 回到原 token 位置。

UVA 不等于免费远程访问。 UVA 只是让 kernel 可以用统一地址访问远端 GPU buffer；真正的通信仍然有高延迟，所以 Comet 才要用 communication thread blocks 隔离远程 I/O。

3.10 Comet 的定位

Comet 的核心不是“把 MoE 拆成几个阶段”，而是：

找到 MoE 中通信和 GEMM 之间共享的 tensor，分析 consumer 在哪个维度上可以独立消费，再按这个维度切分并重排 GroupGEMM tile 顺序，最后用 adaptive thread block assignment 让通信和计算更稳定地重叠。

它比 MiniMax-01 更细，因为它已经进入 GEMM tile / thread block 级别；但它还没有像 FlashMoE 那样把整个 MoE operator 改造成一个 persistent GPU runtime。Comet 更像是在现有 MoE/GEMM 执行栈上，通过 dependency resolving 和 fused kernel scheduling，把 coarse EP overlap 推到 tile-level overlap。

4. FlashMoE：persistent kernel + GPU-resident scheduling

如果说 Comet 是“把 MoE 中通信和 GEMM 之间的 shared tensor 拆到 tile 级，并重排 producer-consumer 顺序”，那 FlashMoE 再往前走了一步：

只优化某两段 pipeline 还不够。只要 MoE 仍然依赖 host-managed scheduling、bulk-synchronous collectives 和大量短 kernel launch，就仍然会有系统级 idle gap。FlashMoE 想把整个 distributed MoE operator 放进一个 GPU-resident persistent kernel 里。

这篇文章适合在读完 Comet 后继续看，因为它们都在讲 fine-grained overlap，但抽象层级不同：

Comet:
  找 shared tensor -> 按 consumer 依赖切分 -> 重排 GEMM tile -> 专门 thread blocks 做通信/计算

FlashMoE:
  把 MoE operator 变成 GPU 常驻 runtime
  用 actor、task、symmetric tensor layout 管理 dispatch / expert compute / combine

4.1 读前概念：GPU kernel、megakernel、kernel launch 和 SM

FlashMoE 比 Comet 难读，很大一部分原因是它不只是讲 MoE，而是在讲 GPU runtime。先把几个底层词拆开。

GPU kernel 是运行在 GPU 上的函数。CPU 端调用 CUDA runtime，把一个 kernel 发射到 GPU 上执行，这个动作叫 kernel launch。一次 launch 会指定 grid / block / thread 结构，例如：

moe_kernel<<<num_blocks, threads_per_block>>>(args);

这个调用发生在 CPU host 侧。GPU 收到任务后，把很多 thread blocks 分配到各个 SM (Streaming Multiprocessor) 上执行。SM 可以理解成 GPU 里的计算工厂，一个 GPU 有很多个 SM；每个 SM 同时驻留若干 thread blocks，内部以 warp 为调度单位执行指令。

kernel launch overhead 指的是：CPU 发起 kernel、CUDA runtime 排队、GPU 接收任务、kernel 之间同步和数据交接带来的额外时间。单个 launch 的开销看起来可能不大，但 MoE 层如果由很多短 kernel 和 collective 串起来，launch gap 会变得明显。

megakernel 通常指把原本多个 GPU kernels / operators 合并成一个更大的 kernel，让这个 kernel 内部自己完成多个阶段的工作。FlashMoE 的 persistent kernel 可以看成一种 MoE megakernel：它不是只算一个 GEMM，而是在一个常驻 kernel 里调度 dispatch、expert FFN、combine 等任务。

这里的“减少 CPU launch”不是说 kernel 可以像 Python 函数那样反复复用，而是：

传统做法:
  CPU launch dispatch kernel
  CPU launch GEMM kernel
  CPU launch activation kernel
  CPU launch GEMM kernel
  CPU launch combine kernel

FlashMoE:
  CPU launch one persistent kernel
  persistent kernel 内部自己调度 MoE tasks

论文 Table 1 里说 FlashMoE 的 #GPU Ops 是 1，而 Comet / Megatron / DeepEP 是几十到几百个。这里的 GPU op 可以理解成需要 host/runtime 调度的一次 GPU operator 或 kernel-level operation。其他框架不是“每个 expert 必然启动一个 kernel”这么简单，实际可能是 group GEMM、通信 collective、elementwise、permute、combine 等很多 operators；但它们仍然需要多次 kernel launch 或 runtime scheduling。FlashMoE 的关键就是把这些阶段压进一个 persistent execution context。

这也是为什么如果不理解 kernel launch，很难感受到 FlashMoE 的收益。它省的不只是某个 GEMM 的计算时间，而是整条 MoE operator 里反复出现的：

host launch -> GPU 执行 -> 等待/同步 -> host launch 下一个 op

4.2 几个底层库和机制：NVSHMEM、DMA、in-device BLAS、CUTLASS

FlashMoE 还频繁提到几个系统名词。

名词 怎么理解
NVSHMEM NVIDIA 的 GPU-side PGAS 通信库。它让多个 GPU 拥有一个对称的 shared address space，并支持 GPU kernel 内部发起 put/get/signal 等 one-sided communication。
DMA Direct Memory Access，直接内存访问。意思是数据搬运可以由硬件/设备引擎完成，不必由 CPU 一字节一字节参与。FlashMoE 语境下重点是 device-initiated remote memory transfer。
in-device BLAS 在 GPU device 侧调用/实现的线性代数操作。BLAS 是矩阵乘、向量运算等基础线性代数接口，FlashMoE 需要在 persistent kernel 内部做类似 GEMM/FFN 的计算。
CUTLASS NVIDIA 开源的 CUDA C++ 模板库，用来写高性能 GEMM / convolution 等 kernels。可以把它看成构建自定义矩阵乘 kernel 的积木。

如果用一句话连接这些概念：

FlashMoE 用 NVSHMEM / UVA / DMA 处理跨 GPU tile 搬运，用 CUTLASS 风格的 device-side GEMM 处理 expert FFN，并把它们组织进一个 persistent MoE megakernel。

4.3 从 Comet 视角看 FlashMoE 的动机

FlashMoE 论文的 Introduction 和 Motivation 里强调三个传统 MoE 系统问题。

FlashMoE overview

第一是 synchronous communication。传统 MoE 常用 AllToAll / AllGather。这些 collective 是 bulk-synchronous 的：参与通信的 GPU 都要进入同一个集体操作，慢的 GPU 会拖住快的 GPU。MoE router 又是动态的，每个 step 分到各 expert 的 token 数不同，所以 straggler 很常见。

第二是 kernel launch overhead。一个 MoE forward 可能包含：

router kernel
dispatch communication kernel
expert GEMM kernel
activation kernel
expert GEMM kernel
combine communication kernel

如果每个阶段都要 host 端发起 kernel，并在 kernel 之间交接，短 kernel 和同步点会堆出很明显的 CUDA API / launch gap。FlashMoE 论文里强调，它的目标是一个 persistent kernel，而不是一串短命 kernel。

第三是 task locality 没有被充分利用。Comet 已经告诉我们：数据到达和 GEMM tile ready 是更细的粒度。但如果调度仍然由 CPU 或粗粒度 collective 管，GPU 很难做到“这个 tile ready 了就马上安排一个 block 去算/传”。

所以 FlashMoE 的目标不是再提出一种新的 router，也不是改变 MoE 数学，而是把 MoE 的执行形态从：

host launches many kernels + collective barriers

变成：

one persistent GPU kernel + device-side tasks + one-sided communication

4.4 MoE 数学没有变：变的是执行系统

FlashMoE 仍然处理普通的 MoE FFN。对每个 token $x_i$，gate 选择 top-k experts，并给出 combine weights。一个简化的 MoE 输出可以写成：

\[y_i = \sum_{j=1}^{k} w_{i,e_j} \cdot E_{e_j}(x_i)\]

其中 $E_{e_j}$ 是第 $j$ 个被选中的 expert，$w_{i,e_j}$ 是 gate 给出的权重。

每个 expert 本身还是普通 FFN：

\[E_e(x)=W_{e,2}\,\sigma(W_{e,1}x+b_{e,1})+b_{e,2}\]

所以要记住：

FlashMoE 改的是 MoE operator 的 runtime，不是 MoE 的函数形式。

这和 Comet 一样：Comet 的 shared tensor decomposition 也不改变数学结果，只改变数据到达、GEMM tile 和 consumer operator 的执行顺序。

4.5 FlashMoE 的核心架构：Figure 5 和 single persistent kernel

项目页里的架构图和论文 Figure 5 很适合作为第一张地图。

FlashMoE architecture

FlashMoE 的 persistent kernel 在每张 GPU 上长期运行。它内部不是所有 thread blocks 都做同一件事，而是用 actor model 把 blocks / warps 分成几类角色：

角色 作用
Processor 执行真正的计算和通信任务，例如 GEMM、element-wise、combine、tile transfer
Subscriber 接收 peer GPU 发来的 tile packets，并解码成 task descriptors
Scheduler 维护 ready queue，把已经 ready 的 task 分配给 Processor

论文实现里，大部分 thread blocks 是 Processor。最后一个 block 被当作类似 “OS block” 的管理块，其中三条 warps 做 Subscriber，一条 warp 做 Scheduler。这个细节不太显眼，但很重要：FlashMoE 并不是让所有 blocks 都参与调度，而是用很少的 GPU 资源做管理，把大部分 SM 留给计算任务。

这句话可以拆成一个更具体的执行画面：

GPU 上有很多 SM。
persistent kernel 启动后，许多 Processor blocks 驻留在这些 SM 上。
Processor blocks 不固定属于某个 expert，也不固定只做 dispatch 或 GEMM。

OS block:
  Subscriber warps: 监听其他 GPU 发来的 messages / packets。
  Scheduler warp:   把已经 ready 的 task 放进队列，并分配给空闲 Processor。

Processor block:
  从队列拿 task。
  task ready 就执行，不 ready 就换下一个。

所以原文说 “ensuring that no GPU SM remains idle throughout the lifetime of the MoE operator” 的意思不是物理上 100% 没有任何空泡，而是：FlashMoE 不希望某些 SM 因为等某个大 collective 或某个固定 expert 而长期闲着。只要队列里还有 ready tasks，Scheduler 就可以把它们分给空闲 Processor blocks。

传统做法更像：

所有 GPU 等 dispatch 完成
然后所有 GPU 做 GEMM
然后所有 GPU 等 combine

FlashMoE 更像：

哪个 tile 到了，就生成哪个 task。
哪个 Processor 空了，就拿一个 ready task。

这就是 Figure 5 的核心：它把 MoE operator 从“固定阶段流水线”变成了“GPU 内部任务系统”。

这和 Comet 的 thread block specialization 有继承关系：

Comet:
  compute blocks + communication blocks

FlashMoE:
  Processor blocks + Subscriber/Scheduler management warps

Comet 的重点是隔离 remote I/O 和 GEMM；FlashMoE 的重点是进一步把 task 的产生、通知、调度、执行都放进 GPU kernel 内部。

4.6 Actor model of concurrent computation 是什么

Actor model 是一种并发计算范式，不是机器学习模型。它最早来自并发系统设计，核心思想是：系统由很多 actor 组成，每个 actor 有自己的状态，通过 message 通信，收到 message 后执行动作、修改状态、继续发送 message。

放到 FlashMoE 里，可以这样对应：

Processor 是执行 actor:
  收到 task descriptor 后执行 tile compute / tile transfer / combine。

Subscriber 是消息入口 actor:
  接收来自其他 GPU 的 messages，并把它们解码成本地可调度的 tasks。

Scheduler 是调度 actor:
  维护 ready tasks，把任务分派给 Processor。

这和普通 CUDA kernel 的区别是：普通 kernel 通常是 “每个 thread/block 按固定索引算自己的那一块”；FlashMoE 的 actor-style kernel 则是 “block 从任务队列拿活干”。它更接近 GPU 上的小型 runtime。

论文 Figure 6 里的 actor dependencies 可以按这条链理解：

remote GPU sends message
-> Subscriber decodes message into task
-> Scheduler observes task readiness
-> Processor executes task
-> Processor may produce another message/task

这个图的重点不是 actor 名字本身，而是 任务依赖不再由 CPU host 串行推进，而是在 GPU 内部通过 message 和 ready queue 推进。

4.7 Unified task abstraction：tile 是 runtime 的基本工作单元

Comet 里我们已经反复说过，tile 不是 token chunk，而是 GEMM / tensor 的静态分块。FlashMoE 沿用这个视角，但把 tile 进一步封装成 runtime task。

论文把 task descriptor 理解成一组 metadata + operator。一个 task 至少要告诉 Processor：

我要处理哪块 tile
这个 tile 属于哪个 expert / 哪张 GPU / 哪个通信阶段
要执行的是 FFN、combine，还是 tile transfer
输入地址和输出地址在哪里
依赖是否已经 ready

用一个很简化的伪代码表达：

while persistent_kernel_is_running:
    task = scheduler.pop_ready_task()

    if task.type == "dispatch_tile":
        processor.transfer_tile(task)

    elif task.type == "ffn_tile":
        processor.gemm_and_activation(task)

    elif task.type == "combine_tile":
        processor.weighted_accumulate_and_writeback(task)

这里的关键不是伪代码本身，而是 ready task 这个概念。传统 MoE 是阶段式：

所有 dispatch 完成 -> 所有 expert GEMM 开始 -> 所有 combine 开始

FlashMoE 是任务式：

某个 dispatch tile 到了 -> 生成 FFN task
某个 FFN output tile 完成 -> 生成 combine task
某个 remote packet 到了 -> Subscriber 解码成 task

这就把 Comet 的“tile ready 就尽早消费”推广成一个 GPU 内部的任务系统。

所谓 Unified task abstraction，就是把 MoE 中看起来不同的事情统一成同一种 task 格式。

比如传统看法里，下面几件事是不同 operators：

dispatch token tile
run expert FFN tile
apply activation
combine output tile
write back remote output

FlashMoE 会尽量把它们都描述成：

task = {
  operator_type,
  input_pointer,
  output_pointer,
  expert_id,
  source_rank,
  target_rank,
  tile_coordinate,
  dependency_state
}

这样 Scheduler 不需要理解“这是 MoE 第几阶段的大 op”，它只需要判断：这个 task 的输入 ready 了吗？ready 就交给 Processor。

这个抽象的好处是动态性。MoE router 让每个 expert 的 token 数动态变化，传统固定 pipeline 很容易某些阶段等人；task abstraction 让运行时只关心 ready tasks，而不是强行按固定阶段推进。

4.8 Tile dimensions：为什么不是越大越好

FlashMoE 论文里有一个很值得记的工程细节：它选择的 tile size 是 $(128,64)$。

这和我们讨论 Comet 时常说的 $128\times128$ 不矛盾。不同 kernel、不同 operator、不同寄存器压力下，最优 tile shape 会变。FlashMoE 给出的直觉是：

  • tile 太小：每个 task 的计算量太少，GPU 利用率低，调度开销相对变大。
  • tile width 太大：每个线程需要保存更多中间值，寄存器使用量上升，可能触发 register spill。
  • tile height 太大：如果 thread 数不变，每个 thread 负责的元素更多，单线程工作量增加，整体并行度下降。
  • thread block 太大：同一个 SM 上能同时驻留的 blocks 变少，SM occupancy 降低。
  • block 内线程越多，同步点 __syncthreads() 的成本也更高，因为更多线程要一起等到 barrier。

这里的 register spill 是理解这段话的关键。GPU 上每个 thread 有很快的寄存器，但寄存器数量有限。如果一个 tile width 太大，线程需要同时保存更多 accumulator / pointer / metadata，中间变量装不下，就会 spill 到 local memory。local memory 名字听起来像本地，其实通常落到显存路径，慢很多。

SM occupancy 指一个 SM 上同时驻留多少 warps / thread blocks。occupancy 高不一定总是更快，但太低会让 GPU 难以隐藏访存延迟。假设 thread block 太大，一个 SM 只能驻留很少 blocks；某个 block 等内存或 barrier 时，SM 没有足够其他 warps 可切换，就容易空转。

所以 FlashMoE 选择 $(128,64)$ 不是因为它在数学上最自然，而是一个平衡：

tile 足够大:
  每个 task 有足够计算量，可以摊薄调度成本。

tile 不太大:
  register pressure 不爆，occupancy 不掉太多，同步成本可控。

所以 tile size 不是“expert 能处理的最大 token 数”，而是 kernel/hardware co-design 的结果。这个点和 Comet 的 granularity mismatch 可以接上：tile 是 GPU 计算和调度的单位，不是 MoE 逻辑层的 token chunk。

4.9 One-sided communication：UVA、NVSHMEM 和 device-initiated DMA

FlashMoE 和 Comet 都想摆脱纯 collective 的粗粒度等待，但 FlashMoE 更强调 one-sided, device-initiated communication。

传统 AllToAll 是：

所有 GPU 进入 collective
runtime 统一搬数据
所有参与方等待完成

FlashMoE 使用 NVSHMEM 建立跨 GPU 的 global address space，并在可用时利用 UVA 做 DMA / RDMA 风格的数据搬运。直觉上，GPU kernel 里的 Processor 或 communication task 可以直接把某个 tile 写到目标 GPU 的某个地址，而不是等 host 发起一个大 collective。

但这里有两个容易误解的点。

第一，one-sided 不等于没有通信成本。远端读写仍然有延迟和带宽限制，只是它不再要求所有 GPU 同步进入同一个 collective。

第二，one-sided 需要非常小心的内存布局。如果两个 GPU 同时往目标 GPU 的同一块 buffer 写，就会发生 write-write conflict。FlashMoE 的 symmetric tensor layout 就是为了解决这个问题。

4.10 Symmetric Tensor Layout：Figure 7 为什么符号很多

FlashMoE 的 Figure 7 讲的是 symmetric tensor layout。它是全篇一个不太容易一眼看懂、但非常关键的设计。

FlashMoE Figure 7(a): Symmetric tensor layout

FlashMoE Figure 7(b): DMA/RDMA state machine

论文里可以把这个 layout 粗略读成：

\[L \sim [P, r, 2, s, n, C_{\mathrm{up}}, h]\]

这里不是逐字符复刻论文公式，而是帮助理解每个维度的含义：

符号 含义 为什么需要
$P$ expert-parallel world size 要区分来自哪个 peer GPU / 写向哪个 peer GPU
$r$ communication rounds MoE forward 里可能有 dispatch round、combine round，不同轮不能覆盖
$2$ incoming / outgoing 两个方向 发送出去的数据和接收进来的数据要分开
$s$ staging buffers temporal buffering，让不同时间段的 writes 落到不同 slot
$n$ 每张 GPU 上 local experts 数 一个 GPU 上可能有多个 experts
$C_{\mathrm{up}}$ upscaled expert capacity 给每个 expert 的 token capacity 预留空间
$h$ token hidden dimension 每个 token vector 的 hidden size

为什么要这么复杂？因为 FlashMoE 要支持 fully non-blocking one-sided writes。不同 GPU 的 task 会同时向 symmetric layout 中写 tile。如果只用一个普通 buffer，就很难避免冲突；要么加锁，要么同步，要么冒险覆盖。

FlashMoE 的思路是加 temporal dimensions：

不同通信轮次用不同 slots
outgoing 和 incoming 分开
不同 staging buffer 分开
不同 peer/local expert/capacity slot 分开

这样一个 remote write 的目标地址由 source GPU、target GPU、round、direction、expert slot、capacity slot 等 metadata 唯一决定。论文还给了 theorem：这个布局是 write-write conflict-free。

如果觉得符号多，可以只抓住 Figure 7 的三个问题：

1. 谁写？source peer rank / direction。
2. 写到哪个阶段？communication round / staging buffer。
3. 写给谁算？local expert id / capacity slot / hidden vector。

可以把 Figure 7(a) 想成一个多维停车场：

第 1 层: 这是 dispatch 轮还是 combine 轮？
第 2 层: 这是 outgoing slot 还是 incoming slot？
第 3 层: 这是第几个 staging buffer？
第 4 层: 这是写给哪个 local expert？
第 5 层: 这是这个 expert capacity 里的第几个 token slot？
第 6 层: 这个 token 的 hidden vector。

远程 GPU 写入时，不是说“随便找一个地方写”，而是根据 routing metadata 算出唯一坐标：

target_address =
  L[peer_rank, round, direction, stage, local_expert, capacity_slot, hidden_slice]

Figure 7(b) 的状态机则是在说明 DMA / RDMA 两种路径都遵循“写数据 + 发 signal + 接收端消费”的顺序。Subscriber 看到 signal 后，才会把对应 packet 解码成 task descriptor，再通知 Scheduler。这样可以避免 Processor 在数据还没写完整时就去读。

Figure 7 复杂，是因为 FlashMoE 要让很多 GPU 同时做 one-sided writes。如果布局少一个维度，就可能出现两个 writer 写到同一个地址，或者 dispatch 和 combine 的数据互相覆盖。它不是为了让 tensor 看起来复杂，而是为了让地址计算天然避免冲突。

这个设计和 Comet 的 shared tensor 思路有微妙区别：

Comet shared tensor:
  关注 producer 和 consumer 怎么共享一块中间 tensor，并按依赖切分。

FlashMoE symmetric tensor:
  关注跨 GPU one-sided writes 怎样在不加同步的情况下安全落位。

换句话说，Comet 主要解决“什么时候可以算”；FlashMoE 还要解决“数据可以安全写到哪里”。

4.11 In-place padding：为什么 padding 也会浪费通信

MoE 里常见一个 capacity 概念：每个 expert 最多接收多少 token。为了让 buffer shape 固定，很多实现会把 expert input padding 到 capacity。

例如某个 expert capacity 是 128，但这次只收到 37 个真实 token。传统实现可能会构造一个 128 行的 buffer：

37 real tokens + 91 null tokens

如果这些 null tokens 也被跨 GPU 传输，就浪费网络带宽。FlashMoE 的 payload efficiency 指的就是避免发送这种无意义 payload。

FlashMoE 的 in-place padding 思路是：

只把真实 token 通过网络发送过去
padding 在本地 symmetric tensor buffer 里完成

这样既满足后续 Processor 按固定 tile shape / aligned capacity 读取的需求，又避免把 null token 当作真实数据走网络。

这个点有点“不显眼但很实用”：MoE 的动态路由会导致每个 expert 的真实 token 数远小于或不等于 capacity，如果系统总是按 capacity 传输，那么稀疏计算节省下来的部分收益会被无效通信吃掉。

4.12 FlashMoE 的执行流程：从一个 token tile 的视角看

可以用一个 tile 的生命周期来理解 FlashMoE。

1. Gate / routing 产生 token -> expert 的映射。

2. 对某个目标 expert，Processor 生成 dispatch tile task。

3. 如果 expert 在远端 GPU：
   Processor 通过 NVSHMEM / UVA 把 tile 写入目标 GPU 的 symmetric tensor slot。
   同时发送 signal，通知目标 GPU 有 tile 到达。

4. 目标 GPU 的 Subscriber 收到 packet / signal，
   解码成 task descriptor。

5. Scheduler 把 ready task 分给空闲 Processor。

6. Processor 执行 expert FFN tile：
   Linear-1 -> activation -> Linear-2。

7. 输出 tile ready 后，生成 combine task。

8. combine task 根据 routing weight 做 weighted accumulation，
   如果原 token 在远端，再通过 one-sided write 写回。

这条链路的核心是：每个 tile 的 arrival、compute、combine 都可以独立推进，而不是被整个 MoE layer 的阶段边界卡住。

4.13 和 Comet 的关系：不是替代，而是更底层的 runtime 化

如果只读过 Comet，很容易把 FlashMoE 理解成“另一个更激进的 Comet”。这个说法有一半对，一半不够准确。

相同点：

  • 都认为 MoE 的瓶颈来自通信和计算的粒度不匹配。
  • 都把粒度降到 tile/task 级别。
  • 都不满足于传统 EP overlap 的 token chunk 级 pipeline。
  • 都需要把通信和 GEMM 资源隔离开，避免 remote I/O 破坏计算。

不同点：

维度 Comet FlashMoE
核心抽象 shared tensor dependency resolving GPU-resident actor/task runtime
主要对象 两条 producer-consumer pipeline 整个 distributed MoE operator
通信方式 更细粒度 P2P / kernel 协调 one-sided device-initiated DMA/RDMA
调度位置 fused kernel 内的 block assignment persistent kernel 内 Scheduler 分配 tasks
内存布局重点 shared tensor 的切分方向和执行顺序 symmetric tensor layout 保证 non-blocking writes
解决的额外问题 granularity mismatch kernel launch、collective barrier、payload inefficiency

所以我会这样理解：

Comet:
  把 MoE 的通信-计算 overlap 做到 tile-aware。

FlashMoE:
  把 MoE 的执行环境改成 tile-task runtime。

4.14 读 FlashMoE 时值得记录的“不太显眼”的细节

第一，single kernel 不等于把所有代码机械粘在一个 kernel 里。 真正关键是 persistent kernel 内部有 task abstraction、actor role 和 ready queue。否则只是把阶段写进一个大 kernel，仍然可能内部空转。

第二，Subscriber / Scheduler 是成本，不是免费午餐。 FlashMoE 特意只用最后一个 block 做 OS block，就是为了把管理成本压低。调度能力太弱会供不上 tasks；调度资源太多又会抢走计算 blocks。

第三，tile size 是硬件约束下的折中。 论文选择 $(128,64)$，背后是 register pressure、shared memory、SM occupancy、block synchronization 的平衡，不是理论上越大越好。

第四，symmetric tensor layout 是正确性设计，不只是性能优化。 如果没有这个布局，one-sided writes 可能需要同步或锁；一旦同步变多，FlashMoE 的核心优势就会被抵消。

第五，payload efficiency 是 MoE 特有问题。 因为 router 动态分配 token，capacity padding 很常见。FlashMoE 的 in-place padding 让网络只传真实 token，不传 null token。

第六，实验结果要看清边界。 论文的 evaluation 主要测单个 MoE layer forward，硬件是 8 张 H100；FlashMoE 用 FP32，而很多 baseline 用 FP16，论文认为这反而让 FlashMoE 更吃亏。但这也意味着，读结果时要把精度、forward-only、单层 MoE operator 和完整训练系统区分开。

4.15 FlashMoE 的定位

FlashMoE 可以理解成：

把 distributed MoE layer 从“host 发起的一串 kernels + collectives”改造成“GPU 内部常驻的 tile-task runtime”。

它比 Comet 更系统级、更激进。Comet 解决的是 shared tensor 的依赖切分和 tile 重排；FlashMoE 解决的是整个 MoE operator 如何在 GPU 内部自行调度、通信、计算和写回。

这也是为什么 DeepSeek-V4 那种 expert-wave pipeline 读起来会更像工程折中：它吸收了 fine-grained overlap 的思想，但没有把论文重点放在一个完整 persistent MoE runtime 上。而 FlashMoE 的主张更明确：要突破 MoE 系统瓶颈，不能只看 GEMM，也不能只看通信，必须把 kernel launch、调度、远程写、内存布局和 padding 一起设计。

5. DeepSeek-V4：expert-wave 级别的 MoE pipeline

DeepSeek-V4 的 MoE 优化小节叫 Fine-Grained Communication-Computation Overlap in Expert Parallelism。它借鉴了前面这些工作，但落点更工程化：围绕 DeepSeek-V4 自己的 expert parallelism，把 MoE 层拆成可以流水的 expert waves。

5.1 DeepSeek 的五段流程：不是 MiniMax 那个“五段”

DeepSeek-V4 把 MoE 层执行写成：

Dispatch All-to-All
-> Linear-1 GEMM
-> SwiGLU / FP8 Cast
-> Linear-2 GEMM
-> Combine All-to-All

这里的 Linear-1 和 Linear-2 是 expert FFN 的两个自然 GEMM。

如果 expert 是 SwiGLU 结构，那么 Linear-1 可以理解成同时做 gate/up 两个 projection：

\[\begin{aligned} g_e &= X_e W_{e,\mathrm{gate}} \\ u_e &= X_e W_{e,\mathrm{up}} \\ z_e &= \mathrm{SiLU}(g_e)\odot u_e \\ Y_e &= z_e W_{e,\mathrm{down}} \end{aligned}\]

所以 DeepSeek-V4 Figure 5 里的 Linear-1 对应前两行，SwiGLU / FP8 Cast 对应第三行和量化转换，Linear-2 对应最后一行。它拆的是 expert FFN 执行路径，不是 MiniMax-01 里 a2a -> allgather -> compute -> reduce-scatter -> a2a 那种并行通信路径。

5.2 什么是 wave

一个 wave 可以理解成一小批 experts。例如某个 MoE 层有很多 experts，不必等所有 experts 的 dispatch 都完成后再一起计算，而是：

wave 0 的 token 到了 -> 先算 wave 0 的 experts
wave 1 的 token 正在传 -> 传完后接着算 wave 1
wave 0 算完后 -> 立刻 combine wave 0 的输出

理想化时间线：

t0: dispatch wave 0
t1: compute  wave 0 | dispatch wave 1
t2: combine  wave 0 | compute wave 1 | dispatch wave 2
t3: combine  wave 1 | compute wave 2 | dispatch wave 3

也就是说，DeepSeek-V4 的切分粒度是 expert wave。它既不是 MiniMax-01 的 token group，也不是 Comet 的 shared tensor tile，也不是 FlashMoE 的完整 persistent task runtime。

这里有一个很容易混淆的点：wave 的切分对象更接近 local experts，而不是 token 序列本身。router 会让不同 experts 收到不同数量的 tokens；某个 wave 只要它包含的 local experts 已经拿到足够的输入，就可以开始执行对应 expert GEMM。没有必要等本 rank 上所有 experts 都完成 dispatch。

因此它想解决的是 MoE 里典型的 long-tail 问题：某些 experts 的 token 较多或通信到达更慢，如果整层同步等待，其他已经 ready 的 experts 会空转；如果按 wave 推进，ready 的部分先算，没 ready 的部分继续通信。

5.3 为什么它能隐藏通信

如果计算足够重，通信就可以被计算盖住。DeepSeek-V4 给出的判断可以理解成：

\[\frac{C}{B} \leq \frac{V_{\mathrm{comp}}}{V_{\mathrm{comm}}}\]

其中：

  • $C$：设备峰值计算能力。
  • $B$：通信带宽。
  • $V_{\mathrm{comp}}$：MoE 层计算量。
  • $V_{\mathrm{comm}}$：MoE 层通信量。

如果右边足够大，说明每传一点数据都能对应很多计算，那么通信更容易被计算隐藏。DeepSeek-V4-Pro 里每个 token-expert pair 大约需要 $6hd$ FLOPs，通信量大约是 $3h$ bytes，所以比例大约是：

\[\frac{V_{\mathrm{comp}}}{V_{\mathrm{comm}}} \approx 2d\]

当 $d=3072$ 时，就是约 $6144$ FLOPs/Byte。这说明在合适硬件条件下，MoE expert 计算可以覆盖很大一部分通信。

更具体地说，DeepSeek-V4-Pro 的一个 token-expert pair 需要大约 $6hd$ FLOPs：gate projection、up projection、down projection 各贡献一部分。通信量大约是 $3h$ bytes：dispatch 输入和 combine 输出不完全同精度，论文按 FP8 dispatch 与 BF16 combine 做估算。这个比例说明：只要系统能把通信切细并和 GEMM 对齐，通信就有机会被 Tensor Core 计算遮住。

5.4 公开代码：MegaMoE2 对应什么

论文说已经公开 CUDA-based mega-kernel MegaMoE2，对应 DeepGEMM PR #304。PR 描述里明确写到：它把 Dispatch -> Linear-1 -> SwiGLU -> Linear-2 -> Combine 融合进一个 mega-kernel，并重叠 NVLink communication 与 Tensor Core computation。

这个公开代码给我们几个很重要的判断依据。

第一，DeepSeek 这里不是只写了一个普通 PyTorch MoE。Hugging Face model repo 里的 inference/model.py 更像结构可读版本；真正的 MoE runtime 优化在 DeepGEMM 的 MegaMoE2 kernel 里。

第二，MegaMoE2 的 Python 测试把 fused path 和 legacy baseline 放在一起。legacy baseline 的逻辑可以概括成：

recv_x = dispatch(x, topk_idx, topk_weights)
l1_y = grouped_gemm(recv_x, l1_weights)
l1_y = swiglu_and_cast(l1_y, topk_weights)
l2_y = grouped_gemm(l1_y, l2_weights)
y = combine(l2_y)

而 fused path 变成：

buffer = get_symmetric_buffer(...)
w1, w2 = transform_weights_for_mega_moe(...)
fp8_fp4_mega_moe(
    output,
    w1,
    w2,
    buffer,
    cumulative_local_expert_recv_stats=...,
)

这不是逐行引用源码，而是把公开测试里的结构抽象出来。重点是：MegaMoE2 不再让 host/runtime 一段一段驱动 dispatch、GEMM、activation、GEMM、combine，而是把这些阶段放进一个 fused kernel/runtime 里调度。

第三，Python API 层没有暴露 num_waves 这种模型超参数。wave 更像 kernel 内部的调度粒度：kernel 根据 topk_idx、topk_weights、local expert receive statistics、symmetric buffer 和变换后的 expert weights，决定哪些 experts 已经 ready，哪些结果可以 combine。

5.5 结合 PR #316 看配置和收益

DeepGEMM Benchmark PR #316 给了和 DeepSeek-V4 对齐的测试配置：

模型 Experts Top-k Hidden size $h$ Intermediate size $d$ EP
DeepSeek-V4-Flash 256 6 4096 2048 8
DeepSeek-V4-Pro 384 6 7168 3072 8

这里的 top-k=6 表示每个 token 会路由到 6 个 experts；EP8 表示 experts 分布在 8 个 expert-parallel ranks 上。benchmark 里 MegaMoE2 相对 legacy baseline 的加速大致在 1.50x 到 1.96x 之间，小 batch / latency-sensitive 场景收益尤其明显。

这也解释了为什么 DeepSeek-V4 论文特别提到 RL rollout 和 high-speed agent serving：这些场景 batch 未必很大，但 latency 很敏感；如果整层 MoE 等待所有 experts 同步，通信 bubble 会更明显。按 wave 推进之后，局部 ready 的 experts 可以先算，上一批结果可以先回传，pipeline 更容易保持饱和。

5.6 论文里几个容易漏掉的系统细节

第一，dispatch 是 pull-based。 DeepSeek-V4 论文说 dispatch 阶段由每个 GPU 主动从远端 GPU 读取 activations，而不是让远端 GPU 细粒度 push 过来。原因是 fine-grained push 需要大量低延迟通知，当前硬件上 signaling overhead 很难忽略；pull-based 让当前 GPU 根据本地需要发起读取，更容易和本地 wave 调度配合。

第二，带宽不是唯一瓶颈。 论文用 $\frac{C}{B}\leq \frac{V_{\mathrm{comp}}}{V_{\mathrm{comm}}}$ 说明只要计算通信比足够高，通信可以被计算遮住。此时继续只提高互联带宽不一定最优，反而要关注 compute、HBM、NVLink 和 power budget 是否能同时撑住。

第三，SwiGLU / cast 也在流水线上。 Figure 5 不是只画了两次 GEMM，中间还有 SwiGLU / FP8 Cast。如果这个 element-wise 阶段太慢，Linear-1 和 Linear-2 之间也会产生 bubble。论文因此提到，未来更低成本的 activation 可能有利于这种 fine-grained overlap。

5.7 DeepSeek 的定位

DeepSeek-V4 不是在论文里提出一个通用 MoE runtime 系统，而是在自己的大模型系统里落地了一个适合 large-scale EP 的 expert-wave pipeline。

它的重点是：

  • 把 expert computation 切成 waves。
  • 让 dispatch、expert GEMM、combine 同时推进。
  • 尽量让 all-to-all 通信隐藏在 GEMM 计算后面。
  • 结合 FP8/FP4、DeepGEMM 等底层优化服务高吞吐训练和推理。

6. 四篇论文的粒度对比

最重要的是不要把所有“overlap”都理解成同一种切法。

工作 切分粒度 主要对象 解决的问题 一句话
MiniMax-01 token group / process group EP、ETP、EDP 通信 大规模训练中的 a2a、allgather、reduce-scatter 开销 group-level overlap
Comet shared tensor / GEMM tile / thread block Dispatch-L1、L2-Combine 的 producer-consumer 链路 通信粒度和 GEMM tile 粒度不匹配 tile-level dependency-aware overlap
FlashMoE tile task / actor / persistent kernel 整个 distributed MoE operator collective barrier 和 kernel launch 开销 GPU-resident task scheduling
DeepSeek-V4 expert wave Dispatch、Linear-1、Activation、Linear-2、Combine expert parallelism 下通信阻塞计算 expert-wave pipeline

再把它们放到一条“越来越细”的轴上：

MiniMax-01
  token group / process group
    ->
Comet
  shared tensor / GEMM tile / thread block
    ->
DeepSeek-V4
  expert wave pipeline for production EP
    ->
FlashMoE
  persistent kernel + GPU-resident task runtime

这条线不是严格的优劣排序，而是抽象层级不同。

7. 一个统一的理解框架

MoE 优化可以按三层看。

7.1 并行策略层

这一层决定 expert 放在哪些 GPU 上，token 怎么发过去。

典型问题：

  • expert 数量如何映射到 GPU？
  • top-k 后 token 分布不均怎么办？
  • EP、ETP、EDP 怎么组合？
  • all-to-all 和 allgather/reduce-scatter 怎么重叠？

MiniMax-01 主要在这一层。

7.2 Kernel / GEMM 调度层

这一层决定矩阵乘和通信怎么交错。

典型问题：

  • token 到了一部分，GEMM 能不能先算？
  • Linear-2 算出一部分，combine 能不能先发？
  • 哪些 thread block 做通信，哪些做计算？
  • group GEMM 的 tile 顺序怎么排？

Comet 和 DeepSeek-V4 主要在这一层，只是切分对象不同。

7.3 Runtime 层

这一层决定整个 MoE operator 是否还依赖外部 collective 和一串 kernel launch。

典型问题：

  • 能不能一个 persistent kernel 管完整个 MoE？
  • GPU 能不能自己发起远端读写？
  • task ready 之后能不能马上调度？
  • 如何避免 barrier 和 straggler？

FlashMoE 主要在这一层。

]]>DeepSeek V4笔记2026-05-08T00:00:00+00:002026-05-08T00:00:00+00:00https://wqh011128.github.io/blog/2026/05/08/DeepSeekV4DeepSeek V4: Towards Highly Efficient Million-Token Context Intelligence

Link:

[toc]

Main Contributions

  • 1M Context: DeepSeek-V4 系列把上下文长度推到 1M tokens。重点不是单纯把 RoPE 拉长，而是围绕长上下文重新设计 attention、KV cache 和 sparse selection。
  • Hybrid Attention Architecture: 使用 Compressed Sparse Attention (CSA) 和 Heavily Compressed Attention (HCA) 混合结构，把近邻 token 精确保留，把远程 token 压缩后选择性读取。
  • Manifold-Constrained Hyper-Connections (mHC): 把传统 residual stream 扩成多条 hidden streams，并用受约束的 mixing 矩阵进行合并与传播。
  • MoE Routing Update: MoE router 支持 sqrt softplus 分数函数和前若干层的 hash routing。
  • MTP: 继续使用 Multi-Token Prediction 思路，为训练提供 next-n token 信号，也为后续推理辅助留下结构入口。
  • Training / Post-training: 官方模型卡中提到 32T+ tokens 预训练、Muon optimizer，以及 SFT + GRPO + on-policy distillation 的后训练路线。

官方给出的模型规模如下：

Model Total Params Activated Params Context Length Precision
DeepSeek-V4-Flash 284B 13B 1M FP4 + FP8 mixed
DeepSeek-V4-Pro 1.6T 49B 1M FP4 + FP8 mixed
DeepSeek-V4-Flash-Base 284B 13B 1M FP8 mixed
DeepSeek-V4-Pro-Base 1.6T 49B 1M FP8 mixed

DeepSeek V4 performance

1. 整体架构

我先把 DeepSeek-V4 看成三条线：

  1. Attention 线： 用 sliding window 处理最近 token，用 CSA / HCA 处理远程 token。
  2. Residual 线： 用 mHC 替代简单 residual，让 hidden state 在多条 stream 之间混合。
  3. MoE 线： 继续走 sparse MoE，但 router 分数和部分层路由方式有变化。

一个比较接近源码的 block 级流程如下：

flowchart TD
    A[input_ids] --> B[Token Embedding]
    B --> C[Expand to n_hc streams]
    C --> D[mHC pre mix for attention]
    D --> E[RMSNorm]
    E --> F[Hybrid Attention: SWA + CSA/HCA]
    F --> G[mHC post mix]
    G --> H[mHC pre mix for MoE]
    H --> I[RMSNorm]
    I --> J[MoE FFN: hash/top-k router + experts]
    J --> K[mHC post mix]
    K --> L[Hyper Head + RMSNorm + LM Head]

在官方 inference 代码里，这个结构大概对应：

  • Transformer: embedding、主干 blocks、lm head、MTP blocks。
  • Block: attention + MoE，两处都包了一层 hyper-connection。
  • Attention: sliding window cache + optional compressor + optional indexer。
  • MoE: router、routed experts、shared expert。

需要注意：Transformers 版本为了兼容 generate、cache 和 HF API，类名更细；官方 inference 版本更像“论文结构的最小可读实现”。这篇笔记会同时参考两者，但讲概念时优先贴近官方 inference 写法。

2. 和 DeepSeek-V3 / V3.2 的关系

DeepSeek-V3 的两个核心词是 MLA 和 DeepSeekMoE。V4 没有丢掉 MoE 路线，但长上下文这件事变成了中心问题。

可以先这样理解：

模型 主要 attention / cache 思路 MoE 长上下文处理重点
DeepSeek-V2 MLA 降低 KV cache DeepSeekMoE 经济型推理
DeepSeek-V3 MLA + MTP + MoE 训练优化 DeepSeekMoE 训练效率和推理能力
DeepSeek-V4 CSA + HCA hybrid attention MoE + hash/top-k router 1M context 下压缩和稀疏选择

V4 的关键不是“把所有历史 token 都完整记下来”。如果 1M tokens 全量 KV cache 直接进 attention，显存和算力都不现实。V4 的路线是：

  • 最近窗口：保留完整 KV，保证局部上下文精度。
  • 更远上下文：压缩成少量 compressed entries。
  • 需要读远程信息时：用 indexer 从 compressed entries 里选 top-k。

所以 V4 的 attention 更像一个分层记忆系统：近处看原文，远处看摘要，必要时再从摘要索引里挑重点。

3. MTP: Multi-Token Prediction

3.1 MTP 是什么

MTP = Multi-Token Prediction。

普通 causal LM 的训练目标是：

\[p(x_{t+1}\mid x_{\le t})\]

也就是每个位置只预测下一个 token。MTP 的思路是：在主 next-token loss 之外，再让模型预测更远的 token，例如：

\[p(x_{t+2}\mid x_{\le t}),\quad p(x_{t+3}\mid x_{\le t}),\quad \ldots\]

可以写成一个简化目标：

\[\mathcal{L}_{\mathrm{MTP}} = \frac{1}{D}\sum_{k=1}^{D} \mathrm{CE}\big(p_k(x_{t+k}\mid x_{\le t}), x_{t+k}\big)\]

其中 $D$ 是额外预测的步数。

直觉上，next-token prediction 只要求模型把眼前一步做好；MTP 会逼模型在 hidden state 里保留更多“接下来几步可能怎么走”的信息。DeepSeek-V3 技术报告里已经使用过 MTP，V4 继续保留了这个方向。

3.2 DeepSeek-V4 里的 MTPBlock

官方 inference 代码里有 n_mtp_layers 和 MTPBlock。从结构上看，MTPBlock 会拿两类信息：

  • 当前主干模型的 hidden state。
  • shifted input token 的 embedding。

先对齐源码接口。官方 inference 里的 MTPBlock.forward 是：

def forward(self, x, start_pos, input_ids):
    e = self.embed(input_ids)
    e = self.enorm(e)
    x = self.hnorm(x)
    x = self.e_proj(e).unsqueeze(2) + self.h_proj(x)
    x = super().forward(x, start_pos, input_ids)
    logits = self.head(x, self.hc_head_fn, self.hc_head_scale, self.hc_head_base, self.norm)
    return logits

也就是说，源码里的 MTPBlock 不会自己构造 shifted tokens。它只要求调用方传入一段 input_ids，然后把这段 token 的 embedding 和上一层 hidden streams 融合。真正的 shift 逻辑属于训练数据/调用逻辑：第 1 个 MTPBlock 传入 x[t+1]，目标是 x[t+2]；第 2 个 MTPBlock 再传入 x[t+2]，目标是 x[t+3]。

这里的 shifted input token 指的是“向后错位后的真实 token”。假设原始序列是：

x0, x1, x2, x3, x4, x5

普通 next-token prediction 在位置 t=2 使用 x0, x1, x2 的上下文预测 x3。

MTP 想继续预测更远的位置。为了预测 x4，MTPBlock 会额外看到一个 shifted token，也就是 x3 的 embedding。这里不是把 x3 当成模型生成结果，而是在训练时把真实序列右移后喂给 MTP 模块，帮助它构造“预测下一步的下一步”的中间状态。

可以把它想成：

位置 主模型看到的上下文 主 LM head 目标 第 1 个 MTPBlock 额外输入 第 1 个 MTPBlock 目标
t=0 x0 x1 x1 x2
t=1 x0, x1 x2 x2 x3
t=2 x0, x1, x2 x3 x3 x4

所以 shifted_input_ids 不是一个新 token，也不是预测 token，而是训练数据中已经存在的 token，只是相对当前位置向后平移了一格。如果有多个 MTP 层，第 2 个 MTPBlock 会继续用更后面的 shifted token，去预测更远的目标。

然后 MTPBlock 会把 shifted token embedding 和主 hidden state 分别投影、相加，再过一个 block，最后接同一个 head 输出 logits。

用等价伪代码表示：

tokens = [x0, x1, x2, x3, x4, x5]
T = len(tokens)

# 主模型：位置 t 的 hidden 用来预测 x[t+1]
# 有效位置：t = 0, 1, 2, 3, 4
main_hidden = backbone(tokens[:-1])
main_logits = lm_head(main_hidden)
main_loss = 0
for t in range(T - 1):
    main_loss += CE(main_logits[t], tokens[t + 1])

# 第 1 个 MTPBlock：额外输入 tokens[t+1]，目标变成 tokens[t+2]
# 有效位置：t = 0, 1, 2, 3
shifted_1 = tokens[1:]               # x1, x2, x3, x4, x5
target_1 = tokens[2:]                # x2, x3, x4, x5
mtp_hidden_1 = mtp_block_1(main_hidden[:-1], shifted_1[:-1])
mtp_logits_1 = lm_head(mtp_hidden_1)
mtp_loss_1 = 0
for t in range(T - 2):
    mtp_loss_1 += CE(mtp_logits_1[t], target_1[t])

# 如果有第 2 个 MTPBlock：继续往后错一格，目标变成 tokens[t+3]
# 有效位置：t = 0, 1, 2
shifted_2 = tokens[2:]               # x2, x3, x4, x5
target_2 = tokens[3:]                # x3, x4, x5
mtp_hidden_2 = mtp_block_2(mtp_hidden_1[:-1], shifted_2[:-1])
mtp_logits_2 = lm_head(mtp_hidden_2)
mtp_loss_2 = 0
for t in range(T - 3):
    mtp_loss_2 += CE(mtp_logits_2[t], target_2[t])

再把一个 MTPBlock 本身展开看，它内部大概是：

class MTPBlock:
    def forward(prev_hidden_streams, shifted_input_ids):
        shifted_emb = embed(shifted_input_ids)
        token_part = e_proj(norm_embed(shifted_emb))
        hidden_part = h_proj(norm_hidden(prev_hidden_streams))

        mtp_hidden = token_part[:, :, None, :] + hidden_part
        mtp_hidden = transformer_block(mtp_hidden, shifted_input_ids)

        logits = shared_lm_head(hyper_head(mtp_hidden))
        return logits

这里有一个容易混淆的点：MTP 不是把模型一次 forward 变成“直接吐多个 token”。在公开 inference 代码里，主 Transformer.forward 默认仍然返回 next-token logits；MTPBlock 更像是额外结构，让 hidden state 多承担几个未来 token 的预测任务。

4. MoE Router、Softplus 和 Hash Route

4.1 Softplus 是什么

Softplus 的定义是：

\[\mathrm{softplus}(x)=\log(1+e^x)\]

它可以看成 ReLU 的平滑版本：

\[\mathrm{ReLU}(x)=\max(0,x)\]

区别是：

  • ReLU 在 $x<0$ 时直接变成 0。
  • Softplus 始终大于 0，并且在 0 附近是平滑的。

这对 router 分数有用，因为 MoE router 最后要得到非负权重。如果使用 softplus，负 logits 不会被硬切成 0，而是保留一个很小的正值。

4.2 Hash Route 是什么

普通 top-k router 是：

indices = topk(scores, k)

Hash route 则是：

indices = tid2eid[input_ids]

也就是说，某些层的 expert 选择不再由当前 hidden state 动态决定，而是由 token id 查表得到。

这像是把“这个 token 去哪些专家”提前写进一个表里。对应到源码，hash layer 会有一个 tid2eid 参数，它的形状大致是：

\[[\mathrm{vocab\_size},\ \mathrm{num\_experts\_per\_tok}]\]

4.3 为什么要 Hash Route

这部分论文和源码能确认的是：DeepSeek-V4 确实支持前若干层使用 hash routing。至于动机，我的理解要稍微保守一点：

  • 早期层的 token 表示还比较接近 token identity，用 token id 路由有一定合理性。
  • hash routing 可以减少动态 top-k 路由带来的计算和通信波动。
  • 固定路由可能让早期 expert 分配更稳定，降低某些 batch 下的 routing collapse 风险。
  • 但它也会牺牲一部分上下文自适应能力，所以不适合所有层都这么做。

因此我会把 Hash Route 理解成：早期层用更便宜、更稳定的 token-id 路由，后续层再交给 hidden-state-based top-k router。

4.4 Hugging Face Router 代码

Hugging Face / 官方 inference 的 router 可以按这两步读：先算 routing scores，再决定 selected experts。Hash route 改的是 selected experts 的来源，不是把 expert weights 全部变成常数。

logits = linear(hidden_states, router_weight)
scores = sqrt(softplus(logits))

if hash_routing:
    selected_experts = tid2eid[input_ids]
else:
    selected_experts = topk(scores + correction_bias, k)

router_weights = gather(scores, selected_experts)
router_weights = router_weights / router_weights.sum(dim=-1, keepdim=True)

这里 correction_bias 可以影响 top-k 选择，但真正的 router_weights 仍然来自未加 bias 的 scores。也就是说，selection 和 weighting 可以使用略不同的信号。

5. mHC: Manifold-Constrained Hyper-Connections

论文第 2.2 节的写法是：先介绍 Standard Hyper-Connections (HC)，再说明 mHC 相比 HC 多了哪些稳定性约束。这里我也按论文顺序记。

5.1 Standard Hyper-Connections

标准 HC 的第一步是把 residual stream 的“宽度”扩成 $n_{\mathrm{hc}}$ 条。论文里说的是：

\[\mathbb{R}^{d}\rightarrow\mathbb{R}^{n_{\mathrm{hc}}\times d}\]

其中 $d$ 是真实喂给 Transformer layer 的 hidden size，$n_{\mathrm{hc}}$ 是 Hyper-Connection 的 stream 数量。Hugging Face 配置里对应字段叫 hc_mult，DeepSeek-V4 常见设置是：

hc_mult = 4

对第 $l$ 层之前的 residual state，论文定义：

\[\begin{aligned} X_l &= \left[ x_{l,1}; \ldots; x_{l,n_{\mathrm{hc}}} \right]^T &\in \mathbb{R}^{n_{\mathrm{hc}}\times d} \end{aligned}\]

这里每个 $x_{l,i}\in\mathbb{R}^{d}$。所以它不是把 $d$ 切成 $d/n_{\mathrm{hc}}$，而是维护 $n_{\mathrm{hc}}$ 条完整的 $d$ 维 residual streams。

HC 引入三个线性映射：

\[A_l\in\mathbb{R}^{1\times n_{\mathrm{hc}}}\] \[B_l\in\mathbb{R}^{n_{\mathrm{hc}}\times n_{\mathrm{hc}}}\] \[C_l\in\mathbb{R}^{n_{\mathrm{hc}}\times 1}\]

更新公式是论文的式 (1)：

\[\begin{aligned} X_{l+1} &= B_lX_l+C_lF_l(A_lX_l) \end{aligned}\]

逐项看维度：

  • $A_lX_l\in\mathbb{R}^{1\times d}$：把 $n_{\mathrm{hc}}$ 条 stream collapse 成一条，作为 layer input。
  • $F_l(A_lX_l)\in\mathbb{R}^{1\times d}$：第 $l$ 个实际 layer 的输出，仍然是 $d$ 维。
  • $C_lF_l(A_lX_l)\in\mathbb{R}^{n_{\mathrm{hc}}\times d}$：把 layer 输出写回 $n_{\mathrm{hc}}$ 条 stream。
  • $B_lX_l\in\mathbb{R}^{n_{\mathrm{hc}}\times d}$：旧 residual streams 之间的线性混合。

所以 HC 的关键点是：inner layer 仍然只处理 $d$ 维输入输出，额外扩展的是 residual stream 的宽度。 这给模型增加了一条不同于 hidden size 的扩展轴，但因为 $n_{\mathrm{hc}}$ 通常很小，所以开销相对可控。

5.2 mHC 相比 HC 改了什么

论文说标准 HC 在多层堆叠时可能出现 numerical instability。mHC 的核心改动是：不再让 $B_l$ 成为任意矩阵，而是把 $B_l$ 约束到 doubly stochastic matrices 的流形上。

论文定义的流形是 Birkhoff polytope：

\[\begin{aligned} \mathcal{M} &= \Bigl\{ M\in\mathbb{R}^{n_{\mathrm{hc}}\times n_{\mathrm{hc}}} \ \Big| M\mathbf{1}_{n_{\mathrm{hc}}}=\mathbf{1}_{n_{\mathrm{hc}}}, \ \mathbf{1}_{n_{\mathrm{hc}}}^{\top}M=\mathbf{1}_{n_{\mathrm{hc}}}^{\top}, \ M\ge 0 \Bigr\} \end{aligned}\]

然后约束：

\[B_l\in\mathcal{M}\]

这里的含义很具体：

  • $M\mathbf{1}=\mathbf{1}$：每一行和为 1。
  • $\mathbf{1}^{T}M=\mathbf{1}^{T}$：每一列和为 1。
  • $M\ge 0$：所有元素非负。

论文给出的理由是：这个约束会让 $\lVert B_l\rVert_2$ 被 1 bound 住，因此 residual transformation 是 non-expansive 的。换句话说，$B_l$ 混合多条 stream 时不会无限放大信号。并且 $\mathcal{M}$ 对矩阵乘法封闭，所以多层堆叠时仍然能保持稳定性。

此外，论文还约束 $A_l$ 和 $C_l$ 非负且有界，避免不同 stream 之间因为符号相反而发生 signal cancellation。

5.3 Dynamic Parameterization

mHC 里的 $A_l,B_l,C_l$ 不是固定参数，而是动态生成的。论文把它们分成：

  • dynamic component：由当前输入 $X_l$ 生成。
  • static component：输入无关的可学习 bias。

给定：

\[X_l\in\mathbb{R}^{n_{\mathrm{hc}}\times d}\]

先 flatten 并 RMSNorm：

\[\begin{aligned} \hat{X}_l &= \mathrm{RMSNorm}(\mathrm{vec}(X_l)) \\ &\in \mathbb{R}^{1\times n_{\mathrm{hc}}d} \end{aligned}\]

然后生成 unconstrained raw parameters：

\[\begin{aligned} \tilde{A}_l &= \alpha_l^{\mathrm{pre}} \cdot (\hat{X}_lW_l^{\mathrm{pre}}) + S_l^{\mathrm{pre}} \end{aligned}\] \[\begin{aligned} \tilde{B}_l &= \alpha_l^{\mathrm{res}} \cdot \mathrm{Mat}(\hat{X}_lW_l^{\mathrm{res}}) + S_l^{\mathrm{res}} \end{aligned}\] \[\begin{aligned} \tilde{C}_l &= \alpha_l^{\mathrm{post}} \cdot (\hat{X}_lW_l^{\mathrm{post}})^T + S_l^{\mathrm{post}} \end{aligned}\]

对应维度是：

\[\begin{aligned} W_l^{\mathrm{pre}},W_l^{\mathrm{post}} &\in \mathbb{R}^{n_{\mathrm{hc}}d\times n_{\mathrm{hc}}} \end{aligned}\] \[\begin{aligned} W_l^{\mathrm{res}} &\in \mathbb{R}^{n_{\mathrm{hc}}d\times n_{\mathrm{hc}}^2} \end{aligned}\] \[S_l^{\mathrm{pre}}\in\mathbb{R}^{1\times n_{\mathrm{hc}}}, \quad S_l^{\mathrm{res}}\in\mathbb{R}^{n_{\mathrm{hc}}\times n_{\mathrm{hc}}}, \quad S_l^{\mathrm{post}}\in\mathbb{R}^{n_{\mathrm{hc}}\times 1}\]

其中 $\mathrm{Mat}(\cdot)$ 把 $1\times n_{\mathrm{hc}}^2$ 的向量 reshape 成 $n_{\mathrm{hc}}\times n_{\mathrm{hc}}$ 矩阵。$\alpha_l^{\mathrm{pre}},\alpha_l^{\mathrm{res}},\alpha_l^{\mathrm{post}}\in\mathbb{R}$ 是可学习 gating factors，论文说它们初始化为较小值。这样做的直觉是：训练初期更接近 static component，等训练稳定后再逐步放大 input-dependent dynamic component。

5.4 Applying Parameter Constraints

得到 raw parameters 之后，mHC 才施加约束。

对 input mapping 和 output mapping：

\[A_l=\sigma(\tilde{A}_l)\] \[C_l=2\sigma(\tilde{C}_l)\]

这里 $A_l$ 的范围是 $(0,1)$，$C_l$ 的范围是 $(0,2)$。它们都是非负且有界的。

对 residual mapping，先保证正数：

\[M^{(0)}=\exp(\tilde{B}_l)\]

然后用 Sinkhorn-Knopp 反复做列归一化和行归一化：

\[\begin{aligned} M^{(t)} &= T_r\left(T_c\left(M^{(t-1)}\right)\right) \end{aligned}\]

其中 $T_c$ 表示 column normalization，$T_r$ 表示 row normalization。最终：

\[B_l=M^{(t_{\max})}\]

论文中取：

\[t_{\max}=20\]

这回答了之前那个疑问：20 次 Sinkhorn 发生在 $n_{\mathrm{hc}}\times n_{\mathrm{hc}}$ 的小矩阵上。DeepSeek-V4 常见 $n_{\mathrm{hc}}=4$，所以每个 token 的 $B_l$ 是 $4\times4$，不是 $d\times d$，也不是 $(n_{\mathrm{hc}}d)\times(n_{\mathrm{hc}}d)$。

5.5 和 Hugging Face 变量的对应

HF Transformers 里的 DeepseekV4DecoderLayer 有两个 hyper-connection 模块：

  • attn_hc：包住 self-attention。
  • ffn_hc：包住 MoE / FFN。

调用形式可以概括为：

post, comb, collapsed = self.attn_hc(hidden_states)

对应论文符号：

HF 变量 论文符号 维度 作用
hidden_states $X_l$ $[B_{\mathrm{sz}},S,n_{\mathrm{hc}},d]$ 多条 residual streams
pre $A_l$ $[B_{\mathrm{sz}},S,1,n_{\mathrm{hc}}]$ 或 broadcast 等价形态 input mapping
collapsed $A_lX_l$ $[B_{\mathrm{sz}},S,d]$ 真正喂给 attention / FFN 的输入
comb $B_l$ $[B_{\mathrm{sz}},S,n_{\mathrm{hc}},n_{\mathrm{hc}}]$ constrained residual mapping
post $C_l$ $[B_{\mathrm{sz}},S,n_{\mathrm{hc}},1]$ 或 broadcast 等价形态 output mapping

pre 没有返回，是因为它已经在 collapsed = A_lX_l 这一步用掉了。decoder layer 拿到 collapsed 后调用 attention / FFN，再把输出写回：

post, comb, collapsed = attn_hc(hidden_states)
attn_output = self_attn(norm(collapsed))
hidden_states = post[..., None] * attn_output[..., None, :] + comb @ hidden_states

这就是论文公式：

\[\begin{aligned} X_{l+1} &= B_lX_l+C_lF_l(A_lX_l) \end{aligned}\]

这里为了代码 broadcast，post[..., None] * attn_output[..., None, :] 写法看起来和矩阵公式不同，但含义就是 $C_lF_l(A_lX_l)$。

6. CSA: Compressed Sparse Attention

论文第 2.3 节先给 Figure 3，再按 Compressed Key-Value Entries → Lightning Indexer → Shared Key-Value Multi-Query Attention → Grouped Output Projection → Attention Sink 的顺序拆 CSA。这一节也按这个顺序记。

DeepSeek-V4 CSA architecture

Figure 3 里有三条流：

  1. Sliding Window KV Entries: 最近一小段 token 不压缩，直接参加 attention，用来保留局部细粒度信息。
  2. Compressed KV Entries: 历史 KV 先由 token-level compressor 压缩，再由 top-k selector 选一部分进入 attention。
  3. Lightning Indexer: 不直接输出 hidden state，只给 compressed KV entries 打分并选择 top-k。

一句话概括：CSA 不是把所有历史 token 都扔进 attention，而是先把每 $m$ 个 KV 压成一个 entry，再让每个 query 只 attend 到 top-k 个 compressed entries，同时额外保留最近 sliding window 的原始 KV。

6.1 Compressed Key-Value Entries

论文从输入 hidden states 开始：

\[H\in\mathbb{R}^{n\times d}\]

其中 $n$ 是 sequence length，$d$ 是 hidden size。CSA 先计算两组 KV entries：

\[C^a,C^b\in\mathbb{R}^{n\times c}\]

以及两组对应的 compression weights：

\[Z^a,Z^b\in\mathbb{R}^{n\times c}\]

这里 $c$ 是 head dimension。更细地说：

  • $C^a,C^b$ 是两套 original KV entries，后面会被加权压缩成 compressed KV entries。
  • $Z^a,Z^b$ 不是最终 KV，而是和 $C^a,C^b$ 同形状的 compression logits / weights，用来决定窗口内不同 token 的压缩权重。
  • 上标 $a,b$ 表示两组 interleaved / overlapping series，不是 attention head 编号。

对第 $i$ 个 compressed entry，论文把两个长度为 $m$ 的连续片段拼起来看：

\[\begin{aligned} C^b_{m(i-1):mi-1} &\in \mathbb{R}^{m\times c} \\ C^a_{mi:m(i+1)-1} &\in \mathbb{R}^{m\times c} \end{aligned}\]

也就是说，$b$ 组的 $m$ 个 token 在前，$a$ 组的 $m$ 个 token 紧接在后。对 $C_i^{\mathrm{Comp}}$ 来说，它实际聚合了一个长度为 $2m$ 的邻域信息：

\[\left[ m(i-1),\; m(i+1)-1 \right]\]

但输出仍然只有一个 compressed entry。

论文公式可以理解成两步。第一步，先在这 $2m$ 个位置上做 compression softmax：

\[\begin{aligned} \left[ S^a_{mi:m(i+1)-1}; S^b_{m(i-1):mi-1} \right] &= \mathrm{Softmax}_{\mathrm{row}} \left( \left[ Z^a_{mi:m(i+1)-1}; Z^b_{m(i-1):mi-1} \right] + \left[ b^a; b^b \right] \right) \end{aligned}\]

其中：

  • $S^a,S^b$ 是真正用于加权求和的 compression scores。
  • $b^a,b^b\in\mathbb{R}^{m\times c}$ 是 learnable compression biases。
  • $\mathrm{Softmax}_{\mathrm{row}}$ 表示沿 token/window 这一维归一化；对每个 head dimension，都在这 $2m$ 个候选 token 上分配权重。

第二步，用这些 scores 加权两组 KV entries：

\[\begin{aligned} C_i^{\mathrm{Comp}} &= \sum_{j=mi}^{m(i+1)-1} S^a_j\odot C^a_j + \sum_{j=m(i-1)}^{mi-1} S^b_j\odot C^b_j \end{aligned}\]

其中 $\odot$ 是逐维乘法。因为每个 $C^a_j,C^b_j\in\mathbb{R}^{c}$，所以：

\[C_i^{\mathrm{Comp}}\in\mathbb{R}^{c}\]

现在再回答“为什么有 overlap 还能压缩到 $1/m$”：因为输出 index 还是按 $i$ 走，每 $m$ 个 token 产生一个 $C_i^{\mathrm{Comp}}$。overlap 只改变一个 compressed entry 聚合时能看的范围，让它看见 $2m$ 个 token；它不改变输出 stride。于是 entries 数量仍然约为：

\[\frac{n}{m}\]

所以压缩率仍然约为：

\[\frac{1}{m}\]

一个更具体的例子：如果 $m=4$，那么 $C_i^{\mathrm{Comp}}$ 聚合的是前一个 block 的 4 个 token 和当前 block 的 4 个 token，一共 8 个 token 的信息；但它仍然只输出 1 个 $c$ 维 compressed KV entry。

边界条件也要注意：当 $i=0$ 时，$m(i-1):mi-1$ 这一段不存在。论文说这个 undefined 的 segment 会被忽略。实现上通常可以理解为 mask 掉非法位置，让 softmax 不给这些位置分配权重。

6.2 Lightning Indexer for Sparse Selection

压缩后的 $C^{\mathrm{Comp}}$ 仍然很多。例如 1M context、$m=4$ 时，远程 compressed entries 仍然是十万级。CSA 不会把所有 compressed entries 都送进 core attention，而是借用 DSA strategy，先用一个轻量 indexer 做 top-k selection。

这一节最容易混淆的点是：Indexer 不是最终 attention。它只负责从 $C^{\mathrm{Comp}}$ 中选出 $C_t^{\mathrm{SprsComp}}$，然后 core attention 再读取这些 selected entries。

6.2.1 Indexer keys

论文先对历史 hidden states 做一次和 $C^{\mathrm{Comp}}$ 类似的压缩，但这里压缩出来的是 indexer keys，而不是最终参与 attention 的 KV entries：

\[K^{\mathrm{IComp}}\in \mathbb{R}^{\frac{n}{m}\times c^I}\]

其中 $c^I$ 是 indexer head dimension。第 $s$ 个 compressed block 的 indexer key 记作 $K_s^{\mathrm{IComp}}\in \mathbb{R}^{c^I}$。可以把它理解成每个 compressed block 的“检索向量”。真正被选中的内容仍然来自 $C_s^{\mathrm{Comp}}$。

6.2.2 Indexer queries

对 query token $t$，论文从它的 hidden state $h_t$ 生成 indexer queries。第一步是 down-projection：

\[c_t^Q = h_t \cdot W^{DQ}\]

然后再 up-projection 成 $n^{I}_{h}$ 个 indexer query heads：

\[\begin{aligned} \left[q^{I}_{t,1};q^{I}_{t,2};\ldots;q^{I}_{t,n^{I}_{h}}\right] &= q_t^I = c_t^Q \cdot W^{IUQ} \end{aligned}\]

维度对应关系如下：

  • $h_t\in \mathbb{R}^{d}$：query token $t$ 的输入 hidden state。
  • $c_t^Q\in \mathbb{R}^{d_c}$：query compressed latent vector。
  • $W^{DQ}\in \mathbb{R}^{d\times d_c}$：indexer query 的 down-projection matrix。
  • $W^{IUQ}\in \mathbb{R}^{d_c\times c^I n^{I}_{h}}$：indexer query 的 up-projection matrix。
  • $q^{I}_{t,h}\in \mathbb{R}^{c^I}$：第 $h$ 个 indexer query head。

这里的“先 down 再 up”不是为了还原原始 hidden state，而是为了用较低秩的 $c_t^Q$ 生成多头 indexer queries。后面的 core attention query 也会复用这个 $c_t^Q$。

6.2.3 DSA index score

接下来，论文给每个 indexer query head 一个 query-dependent weight：

\[\begin{aligned} \left[w^{I}_{t,1};w^{I}_{t,2};\ldots;w^{I}_{t,n^{I}_{h}}\right] &= w_t^I = h_t \cdot W^w \end{aligned}\]

其中 $W^w\in \mathbb{R}^{d\times n^{I}{h}}$，$w^{I}{t,h}\in \mathbb{R}$。然后计算 query token $t$ 对 preceding compressed block $s$ 的 index score：

\[\begin{aligned} I_{t,s} &= \sum_{h=1}^{n^{I}_{h}} w^{I}_{t,h}\cdot \mathrm{ReLU}\left(q^{I}_{t,h}\cdot K_s^{\mathrm{IComp}}\right) \end{aligned}\]

这里 $s<\left\lfloor t/m\right\rfloor$，表示只从 query token 之前的 compressed blocks 中选。这个公式可以拆成三层看：

  1. $q^{I}_{t,h}\cdot K_s^{\mathrm{IComp}}$：第 $h$ 个 indexer head 对 block $s$ 的相似度。
  2. $\mathrm{ReLU}(\cdot)$：只保留正相关的 block。
  3. $w^{I}_{t,h}$：让不同 query token 动态决定哪些 indexer heads 更重要。

得到所有 $I_{t,:}$ 之后，用 top-k selector 保留对应的 compressed KV entries：

\[\begin{aligned} C_t^{\mathrm{SprsComp}} &= \left\{ C_s^{\mathrm{Comp}} \mid I_{t,s}\in \mathrm{Top}\text{-}k\left(I_{t,:}\right) \right\} \end{aligned}\]

注意 top-k 选的是 $C_s^{\mathrm{Comp}}$，不是 $K_s^{\mathrm{IComp}}$。$K^{\mathrm{IComp}}$ 只负责打分，$C^{\mathrm{Comp}}$ 才是后面 core attention 的 key/value 来源。

6.2.4 伪代码

def lightning_indexer(h_t, K_IComp, C_Comp, top_k):
    # Eq. (13): low-rank latent for the query token
    c_Q_t = h_t @ W_DQ

    # Eq. (14): n_h^I indexer query heads, each with dimension c^I
    q_I_t = (c_Q_t @ W_IUQ).reshape(n_h_I, c_I)

    # Eq. (15): query-dependent weights for indexer heads
    w_I_t = h_t @ W_w

    scores = []
    for s, K_s_IComp in enumerate(K_IComp):
        # Eq. (16): DSA-style sparse selection score
        score = 0.0
        for h in range(n_h_I):
            score += w_I_t[h] * relu(dot(q_I_t[h], K_s_IComp))
        scores.append(score)

    selected_blocks = topk(scores, k=top_k)

    # Eq. (17): selected compressed KV entries for core attention
    return [C_Comp[s] for s in selected_blocks]

6.3 Shared Key-Value Multi-Query Attention

选出 $C_t^{\mathrm{SprsComp}}$ 之后，CSA 才进入真正的 core attention。论文这里采用 Multi-Query Attention 的形式：每个 compressed KV entry 同时作为 key 和 value，而 query 有 $n_h$ 个 heads。

先从前面已经得到的 $c_t^Q$ 生成 core attention queries：

\[\begin{aligned} \left[q_{t,1};q_{t,2};\ldots;q_{t,n_h}\right] &= q_t = c_t^Q \cdot W^{UQ} \end{aligned}\]

其中 $n_h$ 是 query heads 数量，$W^{UQ}\in \mathbb{R}^{d_c\times c n_h}$，每个 $q_{t,i}\in \mathbb{R}^{c}$。论文特别说明：这里的 latent query vector $c_t^Q$ 和 indexer queries 使用的是同一个 $c_t^Q$。

然后对每个 query head 做 core attention：

\[\begin{aligned} o_{t,i} &= \mathrm{CoreAttn} \left( \mathrm{query}=q_{t,i}, \mathrm{key}=C_t^{\mathrm{SprsComp}}, \mathrm{value}=C_t^{\mathrm{SprsComp}} \right) \end{aligned}\]

其中 $o_{t,i}\in \mathbb{R}^{c}$。Figure 3 里还会把 selected compressed KV entries 与 sliding window KV entries concat 后送入 shared key-value MQA；公式 (19) 写的是 sparse compressed branch 的核心形式。保留 sliding window KV 的目的，是让最近 token 的 fine-grained dependencies 不完全依赖压缩表示。

6.4 Grouped Output Projection

Core attention 之后会得到 $n_h$ 个 head outputs：

\[\begin{aligned} \left[o_{t,1};o_{t,2};\ldots;o_{t,n_h}\right] &= o_t \in \mathbb{R}^{c n_h} \end{aligned}\]

如果直接把 $o_t\in \mathbb{R}^{c n_h}$ 投影回 $d$ 维 hidden state，计算量大约是：

\[\mathcal{O}\left(c n_h d\right)\]

论文说 DeepSeek-V4 中 $c n_h$ 很大，所以直接做 output projection 会很贵。Grouped Output Projection 的做法是先把 $n_h$ 个 outputs 分成 $g$ 组。第 $i$ 组输出满足：

\[o^G_{t,i}\in \mathbb{R}^{c\frac{n_h}{g}}\]

每一组先投影到一个更小的 intermediate output：

\[o^{G'}_{t,i}\in \mathbb{R}^{d_g}, \qquad d_g < c\frac{n_h}{g}\]

最后再把所有组的 intermediate outputs 拼接：

\[\begin{aligned} \left[o^{G'}_{t,1};o^{G'}_{t,2};\ldots;o^{G'}_{t,g}\right] &\in \mathbb{R}^{d_g g} \end{aligned}\]

并投影成最终 attention output：

\[\hat{o}_t\in \mathbb{R}^{d}\]

直观计算量可以写成两段：

\[\mathcal{O}\left(c n_h d_g\right) + \mathcal{O}\left(g d_g d\right)\]

因为 $d_g < c\frac{n_h}{g}$，第一段相当于先在每组内部做 bottleneck compression；第二段再做跨组混合。它比直接 $c n_h\rightarrow d$ 更省的核心原因是：不再让一个巨大矩阵同时处理所有 heads 到 $d$ 维，而是先把每组 head outputs 压到较小的 $d_g$。

6.5 Attention Sink

论文在 CSA / HCA 的 core attention 里还用了 Attention Sink。这一节的 $h$ 是 attention head 的编号，不是 mHC 里的 $n_{\mathrm{hc}}$。

普通 attention 会把一个 head 的全部权重都分给真实 token / compressed block：

\[\sum_j s_{h,i,j}=1\]

但在长上下文压缩场景里，有些 compressed block 对当前 query 可能没用。如果没有 sink，这个 head 还是必须把 100% 概率质量硬分给它们，可能引入噪声。

论文的做法是给每个 attention head 加一个可学习的 sink logit $z’_h$，并把 $\exp(z’_h)$ 加进 softmax 分母：

\[\begin{aligned} s_{h,i,j} &= \frac{\exp(z_{h,i,j})} {\sum_k\exp(z_{h,i,k})+\exp(z'_h)} \end{aligned}\]

其中：

  • $z_{h,i,j}$：第 $h$ 个 attention head 中，第 $i$ 个 query token 对第 $j$ 个 preceding token / compressed block 的 attention logit。
  • $s_{h,i,j}$：对应的 attention score。
  • $z’_h$：第 $h$ 个 attention head 的 learnable sink logit。

这样真实 KV 上的总 attention mass 变成：

\[\begin{aligned} \sum_j s_{h,i,j} &= \frac{\sum_j\exp(z_{h,i,j})} {\sum_k\exp(z_{h,i,k})+\exp(z'_h)} \le 1 \end{aligned}\]

剩下的概率质量被 sink 吃掉，不参与 value 聚合。直觉上，这允许某个 head 表达“这次这条 attention branch 没什么可看的”，而不是强迫它在一堆低相关 compressed entries 里硬选。

6.6 Hugging Face 对应代码

HF / 官方 inference 里的 CSA 关键路径可以分成两块：Compressor 负责生成 compressed entries，Indexer 负责从这些 entries 中选 top-k。下面是保留变量含义后的删减版。

# Compressor: C / Z / compression bias
coff = 1 + overlap
ape = Parameter([compress_ratio, coff * head_dim])
wkv = Linear(dim, coff * head_dim)
wgate = Linear(dim, coff * head_dim)

kv = wkv(hidden_states)       # corresponds to C entries
score = wgate(hidden_states)  # corresponds to Z weights

kv = kv.unflatten(1, (-1, compress_ratio))
score = score.unflatten(1, (-1, compress_ratio)) + ape

if overlap:
    kv = overlap_transform(kv, value=0)
    score = overlap_transform(score, value=-inf)

kv = (kv * score.softmax(dim=2)).sum(dim=2)
kv = norm(kv)

当 overlap=True 时，overlap_transform 对应论文里的 $C^b_{m(i-1):mi-1}$ 和 $C^a_{mi:m(i+1)-1}$ 拼接；当 compress_ratio=4 时，这就是 CSA 的 compressed KV 生成方式。

# Indexer: Eq. (13)-(17)
c_Q_t = h_t @ W_DQ
q_I_t = (c_Q_t @ W_IUQ).reshape(n_h_I, c_I)
w_I_t = h_t @ W_w

raw = q_I_t @ K_IComp.transpose(-1, -2)
scores = (relu(raw) * w_I_t[:, None]).sum(dim=0)

selected_indices = topk(scores, k=index_topk)
C_t_SprsComp = gather(C_Comp, selected_indices)

最后 Attention Sink 对应的是每个 head 一个 learnable sink logit。它不会产生新的 value，只改 softmax 分母：

denom = exp(attn_logits).sum(dim=-1, keepdim=True) + exp(sink_logits)
attn_scores = exp(attn_logits) / denom

7. HCA: Heavily Compressed Attention

HCA = Heavily Compressed Attention。

相比 CSA，HCA 更直接：

  • 不做 overlap。
  • 不做 indexer。
  • 用更大的压缩窗口，例如 $m’=128$。
  • 每 $m’$ 个 token 压成一个 compressed entry。

对应公式更接近：

\[\begin{aligned} C^{\mathrm{HCA}}_i &= \sum_{j=0}^{m'-1} \mathrm{softmax}(Z_{i,j}+B_j)\, C_{i,j} \end{aligned}\]

HCA 的定位不是“精确找远程细节”，而是“用很小代价保存远程上下文的粗粒度记忆”。

Layer type Compress rate Overlap Indexer KV 形态
Sliding window attention 0 No No 最近窗口原始 KV
CSA 4 Yes Yes 原始窗口 KV + selected compressed KV
HCA 128 No No 原始窗口 KV + 全部 heavily compressed KV

官方 inference 默认 compress_ratios 类似：

(0, 0, 4, 128, 4, 128, 4, 0)

这说明不同层混合使用无压缩、CSA 和 HCA。这个设计挺合理：不是所有层都需要同一种长程记忆，一些层负责精确，一些层负责便宜地扫远处。

HCA 在代码上仍然复用 Compressor，关键差异是压缩率更大、不开 overlap、也不走 indexer：

compress_ratio = 128
overlap = False

kv = wkv(hidden_states)
score = wgate(hidden_states) + ape

kv = kv.unflatten(1, (-1, compress_ratio))
score = score.unflatten(1, (-1, compress_ratio))

kv = (kv * score.softmax(dim=2)).sum(dim=2)
kv = norm(kv)

所以 HCA 的代码路径比 CSA 少两件事：没有 overlap_transform，也没有 Indexer -> topk -> gather(C_Comp)。

8. Muon Optimizer

DeepSeek-V4 的更新不只在 attention 和 MoE 上，也包括训练优化器。模型卡和论文摘要里都把 Muon optimizer 列为训练侧改动：它服务于更快收敛和更稳定的大规模预训练，而不是推理时的网络结构。

Muon 通常展开为 Momentum Orthogonalized by Newton-Schulz。它的核心不是像 AdamW 那样给每个参数做逐元素自适应缩放，而是把二维权重矩阵的梯度更新看成一个矩阵，再对 momentum update 做近似正交化。

可以用下面的抽象公式理解：

\[\begin{aligned} G_t &= \nabla_W \mathcal{L}(W_t) \\ M_t &= \mu M_{t-1}+G_t \\ \hat{M}_t &= \mathrm{NewtonSchulz}(M_t) \\ W_{t+1} &= W_t-\eta \hat{M}_t \end{aligned}\]

其中：

  • $G_t$：当前 step 对权重矩阵 $W_t$ 的梯度。
  • $M_t$：带 momentum 的更新方向。
  • $\mathrm{NewtonSchulz}(\cdot)$：用 Newton-Schulz iteration 做近似正交化，让更新方向在矩阵谱上更均衡。
  • $\hat{M}_t$：正交化后的更新方向。

直觉上，Muon 希望避免某些奇异方向被更新得过强、另一些方向几乎不动。对大模型里的线性层矩阵来说，这种矩阵级更新比逐元素缩放更贴近“层”的结构。

一个很粗略的 Newton-Schulz 伪代码是：

G = grad(W)
M = momentum * M + G

X = normalize(M)
for _ in range(num_ns_steps):
    A = X @ X.T
    X = a * X + b * (A @ X) + c * (A @ A @ X)

W = W - lr * X

这段伪代码只表达机制，不等价于 DeepSeek-V4 的完整训练实现。实际训练里还要处理参数分组、混合精度、分布式并行和不同参数类型的 optimizer 分配。对这篇笔记来说，最重要的结论是：Muon 是训练稳定性和收敛速度的优化，不改变推理时的 forward graph；CSA、HCA、mHC、MoE router 才是模型结构里的主要变化。

9. MoE Runtime Optimization

DeepSeek-V4 论文里还有一个容易被忽略但很重要的系统优化：Fine-Grained Communication-Computation Overlap in Expert Parallelism。它不是改 MoE 的数学定义，也不是改变 router 选专家的规则，而是在 expert parallelism 下重新安排 MoE 层的执行顺序。

论文把一个 MoE 层拆成通信和计算交替出现的五段：

Dispatch All-to-All
-> Linear-1 GEMM
-> SwiGLU + FP8 Cast
-> Linear-2 GEMM
-> Combine All-to-All

其中 Linear-1 和 Linear-2 是 expert FFN 内部天然存在的两个 GEMM，不是额外引入的新结构。如果 expert 使用 SwiGLU，那么 Linear-1 通常同时产生 gate/up 两路中间表示，SwiGLU + FP8 Cast 把中间激活重新量化后交给 Linear-2。

9.1 wave 到底是什么

DeepSeek-V4 的关键做法是把 local experts 进一步切成多个 expert waves。一个 wave 不是 token chunk，也不是某个 expert 的一个 GEMM tile，而是一小批 experts 及其对应的 routed tokens。

论文里的调度逻辑可以理解成：

wave 0 的 experts 完成 dispatch -> 立刻开始 Linear-1 / Act / Linear-2
wave 1 的 experts 仍在接收 token -> 通信继续推进
wave -1 的 experts 已经算完 -> output 立刻 combine 回原 rank

进入稳态后，同一时刻可以同时发生三件事：

combine previous wave | compute current wave | dispatch next wave

这就是它和普通 EP overlap 的区别。普通做法经常是“整层 dispatch 完成后再算，整层算完后再 combine”；DeepSeek-V4 让每个 wave 准备好就进入计算，不等其他 experts。这样可以减少 long-tail expert 或小 batch 场景下的空转，论文特别提到 RL rollout 和 high-speed agent serving 这类 latency-sensitive 场景会受益。

9.2 为什么通信能被盖住

论文给出的判断不是“带宽越高越好”，而是看计算通信比：

\[\frac{C}{B}\leq \frac{V_{\mathrm{comp}}}{V_{\mathrm{comm}}}\]

其中 $C$ 是设备峰值计算能力，$B$ 是互联带宽，$V_{\mathrm{comp}}$ 是 MoE 计算量，$V_{\mathrm{comm}}$ 是 MoE 通信量。

对 DeepSeek-V4-Pro，一个 token-expert pair 的 SwiGLU expert 大约需要：

\[V_{\mathrm{comp}}\approx 6hd\]

这里的 $h$ 是 hidden size，$d$ 是 intermediate hidden size。三项来自 gate projection、up projection、down projection。通信量约为：

\[V_{\mathrm{comm}}\approx 3h\]

这里对应 FP8 dispatch 和 BF16 combine。因此：

\[\frac{V_{\mathrm{comp}}}{V_{\mathrm{comm}}}\approx 2d\]

当 $d=3072$ 时，阈值约为：

\[2d=6144\ \mathrm{FLOPs/Byte}\]

含义是：如果硬件的 $C/B$ 不超过这个比例，那么通信理论上可以被 expert GEMM 覆盖；继续堆更多互联带宽的收益会下降，反而要保证计算、显存和网络能同时高负载运行。

9.3 公开代码能看到什么

论文说已经开源 CUDA-based mega-kernel MegaMoE2，位置在 DeepGEMM。DeepGEMM PR #304 对这个 kernel 的描述很直接：它把 dispatch / linear 1 / SwiGLU / linear 2 / combine 融合进一个 mega-kernel，并重叠 NVLink communication 和 tensor core computation。

公开测试文件 tests/test_mega_moe.py 能看到两个路径的对照：

# fused path
buffer = get_symm_buffer_for_mega_moe(...)
transformed_l1_weights, transformed_l2_weights = transform_weights_for_mega_moe(...)
fp8_fp4_mega_moe(
    y,
    transformed_l1_weights,
    transformed_l2_weights,
    buffer,
    cumulative_local_expert_recv_stats=...,
)

# legacy baseline
recv_x, handle = ep_buffer.dispatch(...)
l1_y = grouped_fp8_fp4_gemm(recv_x, l1_weights, ...)
l1_y = swiglu_and_cast(l1_y, topk_weights, ...)
l2_y = grouped_fp8_fp4_gemm(l1_y, l2_weights, ...)
y = ep_buffer.combine(l2_y, handle=handle)

这段对照很有用：它说明 MegaMoE2 不是单独优化某一个 GEMM，而是把 EP dispatch、两次 expert GEMM、SwiGLU/cast、combine 都放进同一个 fused runtime 里处理。Python API 层没有暴露 num_waves 这种模型超参数；wave 更像 kernel 内部的调度粒度，由 routed token metadata、local expert receive stats、symmetric buffer 和 transformed expert weights 共同驱动。

PR #316 的 benchmark 也和论文说法一致：在 EP8 下，DeepSeek-V4-Flash 使用 256 experts、top-k=6、hidden size 4096、intermediate size 2048；DeepSeek-V4-Pro 使用 384 experts、top-k=6、hidden size 7168、intermediate size 3072。不同 batch size 下相对 legacy baseline 大约有 1.50x 到 1.96x 的加速，其中小 batch latency 场景收益更明显。

9.4 论文给硬件和 kernel 的几个观察

DeepSeek-V4 这节最后还有几个系统侧观察，值得单独记下来：

  • 不是无限堆带宽。 当 $C/B$ 已经低于 $2d$ 这个阈值后，继续增加互联带宽的边际收益下降；更重要的是让 compute、memory、network 同时保持高负载。
  • power budget 会变成瓶颈。 极端融合会让 Tensor Core、HBM、NVLink 同时繁忙，功耗墙可能比单看 FLOPs 或带宽更早出现。
  • dispatch 采用 pull-based 思路。 论文说 dispatch 阶段由每个 GPU 主动读取远端 activations，避免 fine-grained push 带来的高通知延迟；如果未来硬件有更低延迟的跨 GPU signaling，push 才会更自然。
  • SwiGLU 本身也可能阻塞 pipeline。 因为 SwiGLU 有指数和除法相关开销，论文建议未来可以考虑更低成本的 element-wise activation，减少 post-GEMM 处理对 GEMM pipeline 的打断。

这类优化和 MiniMax-01、Comet、FlashMoE 都有关，但粒度不同：MiniMax-01 更偏 token group / process group 级 overlap；Comet 更偏 shared tensor / GEMM tile 级重排；FlashMoE 更激进，把 MoE 执行放进 persistent kernel 里调度；DeepSeek-V4 则采用 expert-wave pipeline 来服务自己的大规模 MoE 推理与训练系统。

更完整的比较我单独整理在这里：MoE 优化的探索。

10. 总结表

模块 解决的问题 核心机制 代价 源码入口
MTP next-token 信号太局部 预测未来多个 token 额外 MTP block/head 计算 MTPBlock
MoE Router MoE routing scores 需要平滑非负，并选择 experts softplus/sqrtsoftplus scores + top-k/hash route selection 和 weighting 逻辑更复杂 Gate / DeepseekV4TopKRouter
Hash Route 早期层动态路由成本和波动 token id 查表得到 experts 上下文自适应弱一些 tid2eid / DeepseekV4HashRouter
MoE Runtime Optimization expert parallelism 下通信阻塞计算 expert waves + MegaMoE2 fused runtime kernel 调度和实现复杂度更高 DeepGEMM fp8_fp4_mega_moe
mHC 深层 residual 信号传播 多 stream + Sinkhorn mixing 小矩阵归一化开销 DeepseekV4HyperConnection
CSA 远程上下文太长 近邻原始 KV + $C_t^{\mathrm{SprsComp}}$ + shared KV MQA 需要 compressor 和 indexer DeepseekV4CSACompressor
HCA 更便宜的远程记忆 大窗口压缩，无 indexer 细节损失更大 DeepseekV4HCACompressor
Attention Sink compressed KV 无关时不应强制分配满权重 分母加入 learnable sink logit 每个 head 多一个 sink 参数 attention score computation
Lightning Indexer compressed entries 仍然太多 $K^{\mathrm{IComp}}$ 与 $q^I_t$ 计算 $I_{t,s}$ 后 top-k 可能漏选远程关键信息 Indexer / DeepseekV4Indexer
Grouped Output Projection $c n_h\rightarrow d$ 直接投影太贵 先按 $g$ 组压到 $d_g$，再投影到 $d$ 增加 per-group bottleneck attention output projection
Muon Optimizer 大规模预训练收敛和稳定性 momentum update + Newton-Schulz 近似正交化 训练侧额外 optimizer 计算 training optimizer

11. 源码与资料入口

12. 最后的一句话理解

DeepSeek-V4 的主线不是“又把模型做大了”。它更像是在问一个工程问题：当上下文来到 1M tokens，模型到底应该怎样保存、压缩、索引和读取历史信息？

CSA、HCA、Lightning Indexer 解决的是 长上下文的读写成本；mHC 解决的是 深层网络里的信息传播稳定性；MTP、router 和 Muon 则分别服务于 训练信号密度、sparse computation 与预训练优化。把这些合起来，DeepSeek-V4 才有了 1M context 下还能工作的结构基础。

]]>ZERO训练技术2026-05-07T00:00:00+00:002026-05-07T00:00:00+00:00https://wqh011128.github.io/blog/2026/05/07/ZEROZERO训练技术

这篇笔记里我统一使用 ZeRO 这个写法，它指的是 DeepSpeed 提出的 Zero Redundancy Optimizer。理解 ZeRO 最好的方式，不是背“Stage 1/2/3 分别切了什么”，而是顺着一次训练 step 去看：

  1. 前向时每张卡手里有什么。
  2. 反向时梯度怎么产生、怎么通信、怎么留下来。
  3. 参数更新时谁来更、更新完怎么让所有卡重新对齐。

只要把这三件事看顺了，ZeRO-1、ZeRO-2、ZeRO-3 的差异就会非常清楚。

先立一个基线：普通 Data Parallel 怎么训练

在普通的数据并行 Data Parallel, DP 里，假设有 N 张卡，那么每一张卡都会保存完整的一份：

  • model parameters
  • gradients
  • optimizer states

这里的 optimizer states 以 Adam 为例，通常包括：

  • 参数本身对应的 FP16/BF16 权重
  • FP32 master weights
  • 一阶矩 m
  • 二阶矩 v

所以普通 DP 的问题不是“算不动”，而是“每张卡都重复保存一整套训练状态，显存浪费很大”。

一次训练 step 的流程

1. 前向

每张卡都拿着完整模型参数，各自处理不同的 micro-batch。
例如 rank 0 处理 batch 的一部分，rank 1 处理另一部分，但它们前向时用的模型权重是一模一样的完整副本。

2. 反向

每张卡先根据自己的 micro-batch 算出一套本地完整梯度。
注意这时的梯度虽然只来自本卡数据，但张量形状是完整的，因为每张卡前向时持有的就是完整模型。

然后所有卡对这些完整梯度做一次 all-reduce：

  • all-reduce 的作用是：各卡把同一个梯度张量做求和或求平均；
  • 操作结束后，每张卡都会拿到一份完全相同的完整梯度。

3. 参数更新

由于每张卡现在都有：

  • 完整参数
  • 完整梯度
  • 完整 optimizer states

所以每张卡都会独立执行一遍完全一样的 Adam update。
因为输入完全一样，所以更新后的参数也完全一样，不需要额外同步。

普通 DP 的特点

  • 好处：逻辑最直接，实现简单。
  • 代价：参数、梯度、optimizer states 在每张卡上都完整复制，显存冗余最大。

后面的 ZeRO，本质上就是一步一步把这三类状态的“完整复制”改成“按 rank 分片保存”。

一张总对比表先建立全局感觉

方案 参数是否分片 梯度是否分片 Optimizer state 是否分片 前向是否需要 all-gather 反向主要通信 更新后如何恢复一致参数副本 通信开销趋势
普通 DP 否 否 否 否 all-reduce 完整梯度 不需要额外恢复，所有卡本地更新结果天然一致 基线
ZeRO-1 否 否 是 否 all-reduce 完整梯度 更新后的参数分片通过 all-gather 重建完整参数 略高于 DP
ZeRO-2 否 是 是 否 layer-by-layer reduce-scatter 梯度 通常仍需 all-gather 让各卡重新拿到完整参数副本 高于 ZeRO-1
ZeRO-3 是 是 是 是 参数 all-gather + 梯度 reduce-scatter 参数平时就是分片存放，需要时再 all-gather 最高

如果只记一句话：

  • ZeRO-1 切 optimizer states。
  • ZeRO-2 在 ZeRO-1 基础上再切 gradients。
  • ZeRO-3 在 ZeRO-2 基础上再切 parameters。

下面按训练流程逐个展开。

ZeRO-1：只切 optimizer states

ZeRO-1 的核心是：模型参数和梯度仍然完整复制，但 optimizer states 不再在每张卡上都保存完整副本，而是按 rank 分片保存。

先回答最关键的一句

在 ZeRO-1 里：

  • 反向传播结束后，每张卡保留的梯度仍然是完整的；
  • 被切分的重点是 optimizer states，而不是梯度；
  • 每张卡只更新自己负责的参数分片，是因为 optimizer states 的 ownership 已经按参数分片分配给不同 rank 了。

optimizer state 怎么切分

假设参数向量被逻辑上切成 N 份，对应 N 个 rank：

  • rank 0 负责第 0 段参数对应的 optimizer states
  • rank 1 负责第 1 段参数对应的 optimizer states
  • …
  • rank N-1 负责最后一段

这里“负责”指的是：

  • 这段参数的 m
  • 这段参数的 v
  • 这段参数的 FP32 master copy

只保存在对应的 rank 上，而不是每张卡都保存一份。

所以 ZeRO-1 节省的显存主要来自 Adam 状态，因为 Adam 的状态量通常比只看参数本身更大。

一次训练 step 的流程

1. 前向

和普通 DP 完全一样：

  • 每张卡都有完整参数副本；
  • 每张卡用自己的 micro-batch 做前向。

因此前向阶段并不需要额外参数通信。

2. 反向

和普通 DP 也基本一样：

  • 每张卡先得到一份本地完整梯度；
  • 然后对完整梯度做 all-reduce，得到所有卡一致的完整梯度。

所以到反向结束时，每张卡手里仍然有一整份完整梯度，而不是只保留自己的梯度分片。

3. 更新

这一步才是 ZeRO-1 和普通 DP 真正拉开差异的地方。

虽然每张卡都有完整梯度，但并不是每张卡都去更新整套参数，而是：

  1. 每个 rank 只更新自己负责的参数分片。
  2. 这次更新所需的 optimizer states 也只在这个 rank 上存在。
  3. 其他不归自己负责的参数段，本 rank 不做 optimizer update。

换句话说：

  • 梯度是完整可见的；
  • 但 optimizer states 是分片持有的；
  • 所以每个 rank 只对自己那一片参数拥有真正的“更新权限”。

更新后怎么得到统一模型

每个 rank 只更新了自己的参数分片之后，所有卡手里的参数副本其实暂时是不完整的“新状态”。
这时需要把每张卡更新好的参数分片重新汇总，恢复成一致的完整模型副本。

通常用的算子是 all-gather。

all-gather 是什么

可以把 all-gather 理解成：

  • 每张卡拿出自己持有的一段张量；
  • 通信结束后，所有卡都拿到这些分片拼起来的完整张量。

所以在 ZeRO-1 里，all-gather 的作用就是：

  • rank 0 拿出自己更新后的参数分片；
  • rank 1 拿出自己更新后的参数分片；
  • …
  • 最后所有卡都重建出一份一致的完整参数副本。

all-gather 和 all-reduce 的区别

  • all-reduce 是“同位置元素做规约，比如求和/求平均，然后每张卡都得到同一个规约结果”。
  • all-gather 是“每张卡拿出不同分片，最后所有卡都拿到拼接后的完整结果”。

在 ZeRO-1 里：

  • 梯度同步靠 all-reduce
  • 更新后参数重建靠 all-gather

相比普通 DP，ZeRO-1 到底改了什么

相比普通 DP，ZeRO-1 唯一改变的是：
optimizer states 从“每卡完整保存”变成了“按 rank 分片保存”；参数和梯度依然是完整副本。

ZeRO-2：在 ZeRO-1 基础上再切 gradients

ZeRO-2 的核心是：不仅 optimizer states 分片，梯度也分片。
参数在大多数时刻仍然是完整副本，这一点和 ZeRO-1 一样。

最重要的区别是：ZeRO-2 不再让“整网完整梯度”在每张卡上一直保留到反向结束，而是尽量在反向过程中就把梯度规约并切走。

一次训练 step 的流程

1. 前向

前向和 ZeRO-1 一样：

  • 每张卡都持有完整参数副本；
  • 每张卡对自己的 micro-batch 做前向。

因此前向阶段仍然不需要额外参数 all-gather。

2. 反向：重点是 layer-by-layer 的 reduce-scatter

这里是 ZeRO-2 的重点。

在普通 DP 或 ZeRO-1 里，一个直观想法是：

  • 整个模型先反向完；
  • 每张卡手里留着整网完整梯度；
  • 最后统一做梯度同步。

而在 ZeRO-2 里，更像是：

  1. 某一层一旦反向完成，这一层的梯度就已经产生了。
  2. 这时不等全模型其他层全部反向结束，就尽快对这一层梯度做 reduce-scatter。
  3. 规约并切分完成后，本卡只保留自己负责的那一段梯度分片。
  4. 然后继续上一层的反向。

所以它是一个非常典型的 layer-by-layer 梯度规约与切分 的过程。

reduce-scatter 是什么

reduce-scatter 可以理解成两步的融合：

  1. 先做 reduce：把各卡对应梯度做求和或求平均。
  2. 再做 scatter：把规约后的完整结果按分片发给不同 rank。

因此它很像：

all-reduce + scatter

但作为一个融合算子，它避免了先在每张卡上留下完整规约结果、再手动切分出去的中间步骤。

反向后每张卡最终保留什么

这是 ZeRO-2 最重要的结论之一：

  • ZeRO-1 反向结束后，每张卡保留的是完整梯度；
  • ZeRO-2 反向结束后，每张卡保留的是自己负责参数分片对应的那一部分已规约梯度。

也就是说，在 ZeRO-2 里，完整梯度不会像 ZeRO-1 那样长期驻留在每张卡显存里。

这就是它比 ZeRO-1 更省显存的核心原因。

3. 更新

到了更新阶段，ZeRO-2 的逻辑反而比 ZeRO-1 更“对齐”：

  • optimizer states 本来就是按分片存的；
  • 梯度现在也已经按同样的 ownership 分片了；
  • 所以每个 rank 直接拿自己那部分梯度，更新自己那部分参数和 optimizer states。

这一步不再需要“虽然梯度完整，但我只更新其中一部分”的逻辑解释，因为梯度本身也已经被切到对应 rank 了。

4. 更新后参数如何一致

由于参数在训练大部分时候仍按完整副本形式参与前向，所以更新结束后，通常仍需要把各 rank 更新后的参数分片重新同步成一致的完整参数副本，以便下一轮前向继续直接使用完整参数。

这里可以继续理解为依赖参数分片的 all-gather 重建。

为什么 reduce-scatter 比“最后留完整梯度”更省显存

因为它改变了梯度在显存中的驻留方式：

  • 不是整网所有层的完整梯度一直堆在每张卡上；
  • 而是哪一层梯度算出来，就尽快规约并切成分片；
  • 本卡只留下自己需要负责更新的那一部分。

所以显存里长期保留的梯度体积明显下降。

相比 ZeRO-1，ZeRO-2 到底改了什么

相比 ZeRO-1，ZeRO-2 新增的本质变化就是：
梯度不再完整复制，而是在反向过程中按 layer-by-layer 的 reduce-scatter 变成分片梯度。

ZeRO-3：在 ZeRO-2 基础上再切 parameters

ZeRO-3 的核心是：参数、梯度、optimizer states 三者全部分片。

这也是为什么它显存节省最激进，但通信也最重。

参数怎么切分

在 ZeRO-3 里，参数平时不是“每张卡都有完整副本”，而是：

  • rank 0 保存一部分参数
  • rank 1 保存另一部分参数
  • …
  • 每张卡只长期保存自己拥有的参数分片

所以 ZeRO-3 和 ZeRO-2 最大的区别在于：
连前向所需的参数也不再常驻完整副本。

一次训练 step 的流程

1. 前向：按层 all-gather 参数

因为每张卡平时只有参数分片，而前向计算某一层时通常需要该层完整权重，所以流程会变成：

  1. 当前层要做前向之前，各 rank 先把这层参数分片通过 all-gather 拼成临时完整参数。
  2. 每张卡拿着这层完整参数做前向计算。
  3. 这一层算完后，非本 rank 所拥有的那部分参数可以释放掉，只保留本地参数分片。

所以 ZeRO-3 的前向不是“开局就有一整套完整模型”，而是“算哪一层，就临时 gather 哪一层”。

2. 反向：先需要完整参数，再做梯度 reduce-scatter

反向时同样要按层理解。

某一层反向时，需要该层参数参与梯度计算，因此通常也要保证这层的完整参数在计算时可见。
完成该层反向后，这层产生的梯度再像 ZeRO-2 一样走 reduce-scatter：

  1. 该层反向得到梯度。
  2. 对这层梯度做 reduce-scatter。
  3. 每张卡最终只保留自己负责参数分片对应的那部分梯度。

所以 ZeRO-3 的反向同时包含两类事情：

  • 为了算这一层，需要临时拿到这层完整参数；
  • 为了节省梯度显存，梯度出来后又尽快分片规约掉。

3. 更新

到了更新阶段，逻辑和 ZeRO-2 一脉相承：

  • 每个 rank 只持有自己的参数分片；
  • 只持有自己那部分梯度；
  • 只持有自己那部分 optimizer states；
  • 因此只更新自己的参数分片。

ZeRO-3 的通信量是不是变大

答案是：是的，通常会明显变大，而且会更频繁。

原因不是“梯度通信突然变得神秘地更多了”，而是参数本身也进入了按层通信流程：

  • ZeRO-2 主要增加的是梯度 reduce-scatter；
  • ZeRO-3 则在前向和反向两边，都要频繁为当前层参数做 all-gather。

这会带来两个后果：

  1. 通信从更粗粒度的 step 级，同步成了更细粒度的 layer 级。
  2. 参数 gather 和梯度 scatter 都要和计算过程紧密交织。

所以 ZeRO-3 通常对：

  • GPU 间互联带宽
  • 通信与计算重叠 overlap
  • bucket 化和 prefetch

会更加敏感。

为什么说 ZeRO-3 最省显存

因为三类主要训练状态都不再完整复制：

  • parameters 分片
  • gradients 分片
  • optimizer states 分片

这使得每张卡长期持有的 model states 最少。
代价就是：为了把“长期常驻显存”压到最低，必须接受更重的动态通信。

相比 ZeRO-2，ZeRO-3 到底改了什么

相比 ZeRO-2，ZeRO-3 的本质新增变化就是：
参数也从“完整副本常驻”变成了“平时按分片保存，需要某层计算时再临时 all-gather”。

最后用一句话串起来

如果把三种 ZeRO 放成一条连续演进路线，可以这样记：

  1. ZeRO-1：参数和梯度还是完整的，先把 optimizer states 从复制改成分片。
  2. ZeRO-2：在 ZeRO-1 基础上，把梯度也从复制改成分片，关键算子是 layer-by-layer 的 reduce-scatter。
  3. ZeRO-3：在 ZeRO-2 基础上，把参数也从复制改成分片，关键代价是前向和反向都要更频繁地做参数 all-gather。

所以 ZeRO 的本质，不是某一个神奇算子，而是一套越来越激进的状态分片策略：

  • 先切 optimizer states
  • 再切 gradients
  • 最后切 parameters

显存越省，通信越重，这就是 ZeRO-1 到 ZeRO-3 最核心的主线。

补充

这篇笔记先聚焦 ZeRO-1 / ZeRO-2 / ZeRO-3 的训练流程本身。
像 ZeRO-Offload、ZeRO-Infinity、和 tensor parallel / pipeline parallel 的组合方式，这里先不展开，后面如果需要可以单独再开一篇。

]]>强化学习笔记2026-04-23T00:00:00+00:002026-04-23T00:00:00+00:00https://wqh011128.github.io/blog/2026/04/23/RL强化学习初学者速通文档

执行摘要

如果你把这份文档只读一遍，最应该建立的认知有四条。第一，RLHF 不是一个单一算法，而是一条训练管线：先做 SFT，让模型学会“基本回答”；再收集偏好比较数据训练奖励模型；最后用 PPO 之类的策略优化算法，让模型在不偏离参考模型太远的前提下提高“人类更喜欢”的回答概率。OpenAI 的 InstructGPT 就是这条路线的经典代表，而且论文报告称，1.3B 的 InstructGPT 在人工偏好评测中优于 175B 的原始 GPT-3。

第二，PPO 是“稳着改策略”。它本质上是策略梯度方法的工程化强化版：先根据旧策略采样，再用一个“剪切后的替代目标”更新新策略，防止一次更新迈得太大。PPO 很通用，也很成熟，但在 LLM 对齐里通常还要配合参考模型、奖励模型、价值模型、KL 约束和大量训练细节，所以它常被认为“强但重”。

第三，DPO 是“把 RLHF 里的一部分 RL 数学化简掉”。在 Bradley–Terry 偏好模型和 KL 约束设定下，DPO 论文把“奖励模型 + RL”改写成了一个直接作用于策略的分类损失，因此不需要显式奖励模型，也不需要在线采样回路；论文与官方实现都强调其训练更简单、更稳定。对初学者来说，DPO 往往是最容易真正跑起来、也最容易从 loss 曲线和 chosen/rejected 对数概率中建立直觉的方法。

第四，GRPO 是“去掉 critic 的 PPO 亲戚”。它在 DeepSeekMath 中被正式提出，并在 DeepSeek-R1 系列训练中被采用。核心做法是：对同一个 prompt 生成一组回答，用组内相对奖励构造 advantage，而不是单独训练一个价值模型；因此它尤其适合有可验证奖励的任务，如数学、代码、规则格式约束。但你也要知道，GRPO 目前并不是一个像 PPO 那样完全定型、单一版本的算法：原始 DeepSeekMath 公式、DeepSeek-R1 实践和 Hugging Face TRL 当前实现之间，已经出现了对 KL 是否启用、奖励是否按组标准差归一化、损失是否做长度修正等细节分化。初学者看到不同文章写法不一致，不一定是谁错了，更可能是“同一家族的不同工程变体”。

阅读前的最小背景与统一符号

为了把 RLHF、PPO、DPO、GRPO 放在同一个坐标系里，我们先统一最常见的符号。对大语言模型而言，状态通常不再写成经典 RL 的 $s_t$，而更常写成“prompt + 已生成前缀”；策略 $\pi_\theta$ 表示模型在当前上下文下生成下一个 token 的条件概率；$\pi_{\text{ref}}$ 是参考策略，通常来自 SFT 模型；$r$ 是奖励，可能来自奖励模型，也可能来自规则验证器；$V$ 是价值函数，只在 PPO/Actor-Critic 一类方法里需要；$A$ 是 advantage，表示“这个动作比当前平均水平好多少”。这些写法在 PPO、DPO、GRPO 的原始论文与 OpenAI Spinning Up 的策略优化教程中是一致的，只是 LLM 场景把“动作”换成了 token 序列。

符号 含义
$x$ 输入 prompt
$y$ 一条完整回答
$y^+, y^-$ 同一 prompt 下的优选回答 / 弱选回答
$\pi_\theta(y\mid x)$ 当前策略模型
$\pi_{\text{ref}}(y\mid x)$ 参考模型，通常是 SFT 模型
$r_\phi(x,y)$ 奖励模型给出的分数
$V_\psi$ 价值模型
$A$ advantage，表示相对优劣
$\beta$ 控制与参考模型偏离程度的系数
$\epsilon$ PPO/GRPO 中剪切范围
$G$ GRPO 中每个 prompt 采样的候选回答数

把四个概念放进一张图里，初学者最需要记住的是：RLHF 是总流程，PPO 是流程中的一个优化器；DPO 是把这个流程中的一部分目标直接改写；GRPO 则是另一种在线策略优化路线，重点是用“组内相对优势”替代 critic。

flowchart LR
    A[预训练模型] --> B[SFT 监督微调]
    B --> C[采样多个回答]
    C --> D[人类或规则给出偏好/分数]

    D --> E[奖励模型 RM]
    E --> F[PPO 式在线优化]
    B --> F
    F --> G[经典 RLHF 模型]

    D --> H[DPO 直接偏好优化]
    B --> H
    H --> I[DPO 模型]

    D --> J[GRPO 组内相对优势]
    B --> J
    J --> K[GRPO 模型]

上图对应的是行业里最常见的理解方式：经典 RLHF 走“SFT → 偏好比较 → RM → PPO”，DPO 走“SFT → 偏好比较 → 直接优化策略”，GRPO 更偏“SFT/基座模型 → 在线采样多个回答 → 组内相对奖励 → 更新策略”。其中 InstructGPT 和 TL;DR 总结工作展示了经典 RLHF 管线，DPO 论文给出了直接偏好优化的闭式改写，DeepSeekMath 与 DeepSeek-R1 则展示了 GRPO 在推理任务上的用法。

RLHF

通俗直观解释

RLHF 可以理解成“先让模型学会说话，再让人类教它什么叫说得更好”。只做预训练时，模型学到的是“互联网上像人类一样继续写下去”；但这不等于“遵从指令、诚实、无害、对用户有帮助”。因此 OpenAI 在 InstructGPT 里采用了三步法：先做人类演示的 SFT，再做人类比较偏好的奖励模型，最后用 PPO 优化策略。OpenAI 在论文和博客中都把这条路线当作把“语言建模目标”改成“更符合用户意图目标”的关键办法。

这条路线为什么有效？因为在很多真实任务里，“哪个回答更好”比“正确答案唯一是什么”更容易标注。比如总结、对话、写作风格、安全拒答，标注者往往更擅长在两个回答里选一个更好的，而不一定能直接写出最优答案。OpenAI 在 TL;DR 总结工作中收集了大规模摘要比较数据，训练奖励模型，再用 RL 去最大化该奖励，结果是模型在人工偏好评价中明显优于纯监督学习基线，而且跨域迁移到 CNN/DM 新闻总结时也能保持较强效果。

关键数学推导

RLHF 的核心通常分成两层数学对象：一个是偏好模型，一个是带 KL 约束的策略优化目标。最常见的偏好建模方式是 Bradley–Terry 形式：给定同一 prompt $x$ 下两个回答 $y^+$ 和 $y^-$，人类更喜欢 $y^+$ 的概率由两者奖励差决定：

\[P(y^+ \succ y^- \mid x)=\sigma\big(r_\phi(x,y^+) - r_\phi(x,y^-)\big)\]

于是奖励模型的训练就是一个二分类最大似然问题：

\[\mathcal L_{\text{RM}}(\phi)= -\mathbb E_{(x,y^+,y^-)\sim \mathcal D} \left[ \log \sigma\big(r_\phi(x,y^+) - r_\phi(x,y^-)\big) \right]\]

这一步的假设是：人类偏好可以被一个潜在标量奖励函数近似，而且偏好数据主要以成对比较形式出现。这个写法在 DPO 论文回顾 RLHF 管线时写得非常清楚，也与 TL;DR 总结和 InstructGPT 采用的“比较数据 → 奖励模型”路线一致。

有了奖励模型后，经典 RLHF 的策略优化目标通常写成：

\[\max_{\pi_\theta}\; \mathbb E_{x \sim \mathcal D,\; y \sim \pi_\theta(\cdot\mid x)} \big[r_\phi(x,y)\big] \;-\; \beta\, D_{\mathrm{KL}}\!\left(\pi_\theta(\cdot\mid x)\,\|\,\pi_{\text{ref}}(\cdot\mid x)\right)\]

第一项是“追求更高奖励”，第二项是“别离参考模型太远”。初学者最容易漏掉的是第二项的重要性：如果没有 KL 约束，模型很容易为了讨好奖励模型而走到奖励模型并不可靠的区域，出现 reward hacking、模式崩塌或者语言质量劣化。 DPO 论文把这个目标明确写为 prior RLHF 的标准形式；InstructGPT 和后续 TRL/PPO 实现则把它具体化成带参考模型的 PPO 训练。

从工程视角看，RLHF 与普通监督学习最大的不同是：训练目标取决于当前模型的采样结果。这意味着训练是“闭环”的，而不是固定数据上的静态拟合。也正因为如此，PPO 和 GRPO 这类在线方法会比 DPO 更重，但也更灵活。

伪代码

下面这段伪代码概括的是经典 RLHF 管线，而不是任何一家公司的逐字实现。流程结构与 InstructGPT、TL;DR summarization 和 TRL/OpenRLHF 常见实现一致。

# Classical RLHF pipeline

# Step 1: SFT
policy = init_from_pretrained_lm()
policy = supervised_finetune(policy, demonstration_dataset)

# Step 2: Reward Model
preference_pairs = collect_pairs(policy, prompt_dataset)   # (x, y_plus, y_minus)
reward_model = init_reward_model(policy)
for batch in preference_pairs:
    loss_rm = -log_sigmoid(reward_model(x, y_plus) - reward_model(x, y_minus)).mean()
    update(reward_model, loss_rm)

# Step 3: RL optimization
ref_policy = copy(policy)   # frozen reference
value_model = init_value_model()   # PPO-style setups usually need it
for iteration in range(num_iters):
    trajectories = sample_responses(policy, prompt_dataset)
    rewards = reward_model.score(trajectories) - beta * kl(policy, ref_policy)
    advantages = estimate_advantages(rewards, value_model)
    update_policy_with_ppo(policy, value_model, trajectories, advantages)

简短 Python 示例片段

下面的片段只用于说明 TRL 当前 PPOTrainer 需要哪些核心对象。它仍然不是可直接运行的完整脚本，但字段名已经尽量对齐当前文档，避免把 PPOTrainer 误解成只要 policy/ref/reward 三件套就能启动。

# 仅作结构示意：真实训练还需要数据预处理、生成配置、分布式训练等
from transformers import (
    AutoModelForCausalLM,
    AutoModelForSequenceClassification,
    AutoTokenizer,
)
from trl import PPOTrainer, PPOConfig

policy = AutoModelForCausalLM.from_pretrained("your-sft-model")
ref_policy = AutoModelForCausalLM.from_pretrained("your-sft-model")
reward_model = AutoModelForSequenceClassification.from_pretrained("your-rm")
value_model = AutoModelForSequenceClassification.from_pretrained("your-value-model")
tokenizer = AutoTokenizer.from_pretrained("your-sft-model")
train_dataset = load_your_prompt_dataset()

args = PPOConfig(
    learning_rate=3e-6,
    cliprange=0.2,
    kl_coef=0.05,
    num_ppo_epochs=4,
)

trainer = PPOTrainer(
    args=args,
    processing_class=tokenizer,
    model=policy,
    ref_model=ref_policy,
    reward_model=reward_model,
    value_model=value_model,
    train_dataset=train_dataset,
)

# 训练时的闭环仍然是：prompt -> generate -> reward -> advantage/value -> PPO update

实践要点与超参数建议

对初学者最有价值的经验不是“某个神奇超参”，而是先把三阶段的职责分清。SFT 负责把分布收窄到“像样回答”的区域；奖励模型负责把“更好/更差”映射成标量；PPO 负责在 KL 约束下移动策略。如果你的 SFT 很差，后续 RM 和 PPO 往往也不会好，因为比较数据和采样数据都会变脏。OpenAI 的 InstructGPT 和 TRL 的 PPO 文档都把 SFT 作为 RLHF 的前置基础。

在工程细节上，Hugging Face 对 OpenAI 早期 RLHF 代码的复现实验总结出一批非常“接地气”的经验：关闭 dropout、奖励/值函数归一化、学习率退火、奖励白化、adaptive KL、必要时做 rejection sampling，这些都显著影响稳定性；而且他们在复现实验中还特别指出，PPO 训练里 PyTorch/TF 的 Adam 数值行为差异都可能导致更新激进。对初学者来说，这意味着：不要低估实现细节，RLHF 不是“把 PPO 套上去”就完事。

如果你想做一个“能跑通”的入门版本，我的建议是：先不用追求 InstructGPT 规模，优先选择公开框架与小模型。奖励模型可以从序列分类头开始；如果数据很少，优先防止 RM 过拟合；PPO 阶段先严格盯住 objective/rlhf_reward、objective/kl、val/ratio 等指标，因为这些指标能直接告诉你训练是不是“往前走但没走飞”。TRL 文档特别提示，val/ratio 应围绕 1 左右波动，过大或过小都说明相邻策略之间变化过猛。

典型实现与在线文档

就“官方材料”而言，InstructGPT 的完整训练代码公开程度是未指定；但 OpenAI 已公开论文、博客、模型卡/评测仓库，以及更早的 lm-human-preferences 和 summarize-from-feedback 代码资源作为近似参考。对学习者而言，这些足以建立 RLHF 管线认知。

比较值得直接阅读的材料有：OpenAI 的 InstructGPT 论文与博客、OpenAI 的 Learning to Summarize from Human Feedback 论文与仓库、OpenAI 的 lm-human-preferences 代码库、Hugging Face 的 TRL Reward Modeling / PPOTrainer 文档、以及 OpenRLHF 这类偏工程化的高性能开源框架。中文方面，优先看 Hugging Face 中文博客/文档、OpenRLHF 中文文档 和 Spinning Up 中文版，因为它们更接近原始英文材料，而不是二次转述。

常见问题与调试建议

RLHF 最常见的问题，不是 loss 不下降，而是你根本不知道下降的是什么。奖励模型分数升高，不等于人类偏好一定同步升高；KL 很低，不等于模型真的学到了更好的行为；value loss 稳定，也不等于 policy 真在变好。这也是 OpenAI 后续研究 CriticGPT 的原因之一：随着模型越来越强，人类评审越来越难稳定地发现错误，而 RLHF 的标签质量会因此成为瓶颈。

如果你跑 PPO 式 RLHF，先排查四件事。第一，看回答是否很快塌缩成模板化短句；第二，看 KL 是否飞升；第三，看 val/ratio 是否经常远离 1；第四，看 RM 是否在训练集上好、在验证集上差。如果出现这些问题，优先从减小学习率、增大 KL 约束、缩短回复长度、检查 EOS/截断规则、复查 chosen/rejected 数据质量入手，而不是立刻改模型结构。

小结与适用场景

如果你想理解“现代大模型对齐最经典的一条路”，RLHF 必学；如果你想做一个工业上可解释、可扩展、可插入任意奖励函数的系统，RLHF 仍然很重要；但如果你只是想快速从偏好数据起步，RLHF 往往不是最容易的第一站，因为它的工程复杂度明显高于 DPO。这个判断并不是说 RLHF 过时，而是说它更像“全家桶”，适合在你已经会 SFT、懂 PPO、知道 RM 风险之后再系统掌握。

PPO

通俗直观解释

PPO 最容易理解的方式是：每次都朝“更高奖励”的方向走，但只允许走一小步安全步长。如果不加限制，策略梯度会鼓励模型把某些高回报动作的概率一路抬高，结果常常是“一步迈太大”，训练发散。TRPO 用 KL 约束来硬性限制步长，PPO 则用更简单的剪切目标来近似这种“别走太远”的思想，所以它在工程上更易实现。OpenAI 的 PPO 原始论文和 Spinning Up 教程都把它定位为在简单性、样本效率和稳定性之间取得平衡的方法。

放到 LLM 对齐里，PPO 的角色就变成：模型先根据当前策略回答一批 prompt，再根据奖励模型给分；如果某些回答分高，就提高这条回答路径上 token 的概率；但提高的幅度要被剪切项和 KL 项约束住，不然模型会迅速偏离语言质量良好的区域。DeepSeekMath 在介绍 GRPO 时，也先把 PPO 作为“当前 LLM RL 微调中广泛使用的 actor-critic 算法”来对照。

关键数学推导

从策略梯度出发，我们希望最大化策略的期望回报 $J(\theta)$。在旧策略 $\pi_{\theta_{\text{old}}}$ 采样的数据上，可以把新旧策略的差异写成一个比值：

\[r_t(\theta)=\frac{\pi_\theta(a_t\mid s_t)}{\pi_{\theta_{\text{old}}}(a_t\mid s_t)}\]

于是得到经典的替代目标：

\[L^{\text{CPI}}(\theta)=\hat{\mathbb E}_t\big[r_t(\theta)\hat A_t\big]\]

这里 $\hat A_t$ 是 advantage，它表示在状态 $s_t$ 下选动作 $a_t$ 比平均动作好多少。这个目标的问题是：如果直接最大化，策略可能在少数状态上被推得太远。PPO 的关键改写就是加入剪切：

\[L^{\text{CLIP}}(\theta)= \hat{\mathbb E}_t \Big[ \min\big( r_t(\theta)\hat A_t,\; \operatorname{clip}(r_t(\theta),1-\epsilon,1+\epsilon)\hat A_t \big) \Big]\]

这就是 PPO 最核心的公式。论文给出的直觉是：如果 $\hat A_t>0$，那就不希望 $r_t$ 大幅高于 $1+\epsilon$；如果 $\hat A_t<0$，那就不希望 $r_t$ 大幅低于 $1-\epsilon$。这样做的结果不是“完全不让策略变”，而是“超过安全区后，继续变大对目标不再有额外好处”。

在 LLM RLHF 中，PPO 还经常配一个价值函数 $V_\psi$ 来降低方差，并配一个参考模型 KL 惩罚避免策略漂移。DeepSeekMath 对 PPO 的回顾特别指出，PPO 在 LLM 场景里通常需要训练 value function，并在 token 级奖励中加入来自参考模型的 KL 惩罚；这恰好也是后来 GRPO 要“去 critic”的原因。

伪代码

下面的伪代码是 PPO 的最小骨架，结构上与 Schulman 的原始算法和 TRL 的 PPOTrainer 都一致。

initialize policy πθ
initialize value function Vψ

for each iteration:
    rollouts = collect_trajectories(πθ_old)
    rewards = compute_rewards(rollouts)
    advantages = estimate_advantages(rewards, Vψ)

    for epoch in range(K):
        ratio = πθ(a|s) / πθ_old(a|s)
        clipped_ratio = clip(ratio, 1 - eps, 1 + eps)
        policy_loss = -mean(min(ratio * advantages,
                                clipped_ratio * advantages))
        value_loss = mse(Vψ(s), returns)
        loss = policy_loss + c_v * value_loss - c_e * entropy(πθ)
        update(θ, ψ)

简短 Python 示例片段

下面的片段强调的是 PPO 的对象依赖关系：当前策略、参考策略、奖励模型、价值模型、prompt 数据集和 tokenizer/processor 缺一不可。它是结构示意，不是最小可运行脚本。

from transformers import AutoTokenizer
from trl import PPOConfig, PPOTrainer

tokenizer = AutoTokenizer.from_pretrained("your-sft-model")
args = PPOConfig(
    learning_rate=3e-6,
    num_ppo_epochs=4,
    cliprange=0.2,
    vf_coef=0.1,
    kl_coef=0.05,
    gamma=1.0,
    lam=0.95,
)

trainer = PPOTrainer(
    args=args,
    processing_class=tokenizer,
    model=policy_model,
    ref_model=ref_model,
    reward_model=reward_model,
    value_model=value_model,
    train_dataset=prompt_dataset,
)

# prompts -> generate -> reward -> value/advantage -> PPO update

实践要点与超参数建议

对初学者最友好的 PPO 起点，往往不是原论文的通用设置，而是现成的 LLM 实现默认值。以 TRL 为例，PPOConfig 的关键默认值包括：cliprange=0.2、vf_coef=0.1、cliprange_value=0.2、gamma=1.0、lam=0.95、kl_coef=0.05、num_ppo_epochs=4。这组默认值反映了 LLM RLHF 的典型做法：把整段回答的质量看得比传统逐步折扣更重要，因此常把 $\gamma$ 设成 1；同时用 KL 和 value loss 控制训练稳定性。

如果你要调 PPO，优先盯三件事。第一，val/ratio 是否明显脱离 1；第二，objective/kl 是否持续上冲；第三，policy/clipfrac_avg 是否过高。TRL 官方文档直接给了调试建议：val/ratio 应在 1 附近波动，若到 2、1000 或掉到 0.1 这类量级，说明连续两次策略更新差异太大，应先回头检查学习率、奖励尺度、KL 系数和 batch 构造。

再往前走一步，Hugging Face 对 OpenAI 早期 RLHF 的复现材料还给出一批很实用的 PPO 细节：训练中通常关闭 dropout；为 reward 与 policy 做学习率退火；在某些任务中做reward whitening；并通过 adaptive KL 动态调节 KL 系数。这些做法不一定对所有任务都最优，但对“为什么我的 PPO 跑不稳”这个初学者问题非常关键。

典型实现与在线文档

学习 PPO 最推荐的顺序是：先看 PPO 原始论文，理解剪切目标；再看 OpenAI Spinning Up，把公式和实现联系起来；如果你的目标是 LLM 对齐，再看 TRL PPOTrainer 和 OpenAI/复现项目的 RLHF with PPO 实现细节。这样你不会把“机器人控制 PPO”和“语言模型 PPO”混为一谈。

中文学习材料里，Spinning Up 中文版适合建立基础策略优化直觉，Hugging Face 中文博客的 RLHF with PPO 实现细节更适合理解 LLM 对齐中的工程坑。

常见问题与调试建议

PPO 在 LLM 场景最典型的失败模式是“奖励涨了，但文本坏了”。这常见于 KL 太弱、奖励模型可被钻空子、或者 EOS/截断规则处理不当。OpenAI 复现工作中特别提到 rejection sampling、截断、EOS 处理和固定低分惩罚，这些都不是数学主角，但它们会直接改变模型被奖励的输出分布。

第二类问题是“critic 学坏导致 advantage 失真”。如果 value loss 长期震荡、优势估计噪声很大，policy 就会跟着乱动。DeepSeekMath 在提出 GRPO 时正是明确指出：在 LLM 里，value model 常与 policy 模型同量级，内存和算力负担都很大，而且当奖励只在结尾给出时，训练一个逐 token 准确的 value function 还会变得更难。

小结与适用场景

PPO 适合你需要在线探索、奖励函数复杂、且希望保留“显式 RL 闭环”能力的时候。它是最经典、最通用、也最重的一条路。如果你的目标是“先把偏好学习跑起来”，DPO 通常更轻；如果你的目标是“在可验证推理任务里做在线强化学习，但不想背 critic 成本”，GRPO 往往更合适。

DPO

通俗直观解释

DPO 的核心直觉非常适合初学者：既然我们手里已经有“同一 prompt 下哪个回答更好”的配对数据，那为什么还要先学一个奖励模型，再跑一轮 RL，最后才能让策略更偏向好回答？ DPO 的答案是：在特定假设下，可以直接把这个优化目标写成分类损失，于是训练就变成“提高优选回答的相对似然，降低弱选回答的相对似然”。

这也是为什么 DPO 论文会说它“stable, performant, and computationally lightweight”，并强调不需要在微调时从语言模型中在线采样，也不需要显式奖励模型。对于初学者，这意味着：比起 PPO，你可以更快地建立“偏好数据如何改变模型分布”的感觉，因为训练就像一个偏好版的二分类。

关键数学推导

DPO 推导的起点并不是“完全抛弃 RLHF”，而是从 RLHF 的 KL 约束目标出发。假设我们要解的仍是下面这个问题：

\[\max_{\pi}\; \mathbb E_{x\sim \mathcal D,\; y\sim \pi(\cdot|x)}[r(x,y)] -\beta D_{\mathrm{KL}}(\pi(\cdot|x)\|\pi_{\text{ref}}(\cdot|x))\]

DPO 论文说明，这个目标的最优策略可以写成：

\[\pi^*(y\mid x)=\frac{1}{Z(x)}\pi_{\text{ref}}(y\mid x)\exp\!\left(\frac{1}{\beta}r(x,y)\right)\]

其中 $Z(x)$ 是归一化常数。把它取对数并整理，可得：

\[r(x,y)= \beta \log \frac{\pi^*(y\mid x)}{\pi_{\text{ref}}(y\mid x)} +\beta \log Z(x)\]

接下来引入 Bradley–Terry 偏好模型：

\[P(y^+ \succ y^- \mid x)=\sigma\big(r(x,y^+) - r(x,y^-)\big)\]

把上面的 reward 重参数化代进去，$Z(x)$ 会在成对比较中抵消，于是得到只与策略和参考策略有关的偏好概率。最终得到 DPO 损失：

\[\mathcal L_{\text{DPO}}(\theta)= -\mathbb E_{(x,y^+,y^-)} \log \sigma\!\left( \beta \log \frac{\pi_\theta(y^+\mid x)}{\pi_{\text{ref}}(y^+\mid x)} - \beta \log \frac{\pi_\theta(y^-\mid x)}{\pi_{\text{ref}}(y^-\mid x)} \right)\]

这就是 DPO 的核心：把“学奖励再做 RL”改成了“直接学一个更符合偏好的策略”。需要强调的是，这个结论依赖于明确的建模假设，尤其是 KL 约束形式和 Bradley–Terry 偏好建模；因此它不是“所有 RLHF 都被 DPO 完全取代”的数学结论，而是“在这类设定下，目标可被直接改写”的结论。

从实现角度看，DPO 还有一个非常实用的“隐式奖励”视角。TRL 文档把 chosen/rejected 的隐式奖励记成相对参考模型的对数概率比；如果你监控 rewards/chosen、rewards/rejected 和 rewards/margins，其实就在看模型相对参考模型把“好回答”和“差回答”分开得有多明显。

伪代码

下面的伪代码几乎就是 DPO 的本质。

initialize policy πθ from SFT model
freeze reference policy πref

for batch in preference_dataset:   # (x, y_plus, y_minus)
    logp_plus = log πθ(y_plus | x)
    logp_minus = log πθ(y_minus | x)

    ref_logp_plus = log πref(y_plus | x)
    ref_logp_minus = log πref(y_minus | x)

    margin = beta * ((logp_plus - ref_logp_plus) -
                     (logp_minus - ref_logp_minus))

    loss = -mean(log_sigmoid(margin))
    update(θ, loss)

简短 Python 示例片段

下面这段代码非常接近 TRL 官方文档中的最小示例。它足够短，适合你亲自替换模型与数据集跑一个 demo。

from datasets import load_dataset
from trl import DPOTrainer

trainer = DPOTrainer(
    model="Qwen/Qwen3-0.6B",
    train_dataset=load_dataset("trl-lib/ultrafeedback_binarized", split="train"),
)

trainer.train()

实践要点与超参数建议

DPO 的第一原则不是调 $\beta$，而是确保 preference pair 真的可信。如果 chosen / rejected 顺序经常反了，或者两条回答质量几乎没有可比差异，DPO 会学到很奇怪的边界。因为它没有单独的奖励模型做缓冲，所以数据质量问题会更直接地反映到策略里。DPO 论文与 TRL 文档都强调它处理的是“同一 prompt 的优选/弱选完成对”。

$\beta$ 是 DPO 里最值得优先理解的超参数。TRL 文档给出的默认值是 0.1，并把它解释为控制与参考模型偏离程度的关键参数；而后续 $\beta$-DPO 工作又专门指出，$\beta$ 对性能是敏感的。对于初学者，最稳妥的建议是：先用 0.1 跑通，再围绕它做少量扫描，不要一上来就扫特别大范围。

在工程上，TRL 当前默认 disable_dropout=True，并提供 precompute_ref_log_probs=True 来节省显存；此外，padding_free 可以配合 FlashAttention 进一步减少 padding 开销。这些都很适合大模型或长序列场景。

如果你的偏好标签有噪声，可以关注 Robust DPO / EXO 一类扩展。TRL 文档中已经把 label_smoothing 暴露出来，并给出了一些典型值说明，例如 Robust DPO 中常把它解释成标签翻转概率，推荐值可取 0.1 的量级；但如果你只是入门，建议先不要同时引入太多 DPO 变体。

典型实现与在线文档

学习 DPO 的“黄金组合”通常是：先读 原始论文，再看 Stanford 的参考实现仓库，然后直接上 TRL DPOTrainer。这样做的好处是你既知道数学上为什么成立，又知道工程上怎么落地。

中文方面，优先读 Hugging Face 中文博客《从 RLHF 到 DPO》、Hugging Face 中文版 DPOTrainer 文档、以及 使用 DPO 微调 Llama 2 这类带脚本的实战教程。它们跟英文原始材料一致性较高。

常见问题与调试建议

DPO 第一类常见问题是“loss 在降，但模型回答更死板”。这通常意味着模型在过度贴近某类 chosen 模板，或者 $\beta$ / 数据分布让模型相对参考模型的偏移过于单一。此时优先看 rewards/margins、rewards/accuracies、logps/chosen、logps/rejected 是否同步改善，而不是只看总 loss。

第二类问题是“拿未做过 SFT 的底座直接做 DPO”。理论上不是绝对不行，但实践上很容易因为 pair 数据分布太窄而出问题。DPO 论文本身就是在 RLHF 的 SFT 起点上做改写；TRL 文档的示例也默认你已经有可用的 base/SFT 模型。对初学者来说，DPO 最好接在一个已经能正常完成任务的 SFT 模型之后。

小结与适用场景

当你手里有静态偏好数据，又希望用尽可能简单的训练方式完成偏好对齐时，DPO 几乎总是值得优先尝试。它比 PPO 更轻，训练与调试门槛更低，也更适合“先学会偏好优化，再回头学在线 RL”。但它的前提很明确：你要有质量不错的 chosen/rejected 数据，而且最好已经有一个 decent 的 SFT 起点。

GRPO

通俗直观解释

GRPO 可以把它想成“不再问一个回答绝对值多少，而是问它在同组候选里相对表现如何”。对同一个 prompt，你用旧策略一次采样出 $G$ 个回答；然后用奖励模型、规则验证器或准确率函数给它们打分；接着看每个回答相对这组回答的平均水平高多少、低多少；最后用这个“组内相对优势”更新策略。DeepSeekMath 论文把它定义为 PPO 的一个变体，最大卖点是：不再需要单独训练 critic/value model。

为什么这个思路在推理任务上特别流行？因为数学、代码、格式约束这些任务里，很多奖励都可以由规则直接验证，例如最终答案是否正确、输出是否满足模板、测试用例是否通过。DeepSeek-R1 进一步说明，他们在 reasoning 任务上更倾向于使用这类规则型奖励信号，而刻意避免大规模依赖神经奖励模型，因为后者在强 RL 下更容易被 reward hacking。

关键数学推导

GRPO 这部分最容易让人混淆的地方在于：原始 DeepSeekMath 写法、后续 PPO-style 推广写法、以及 TRL 当前实现的工程修正式，并不是同一个公式。 把它们拆开看会清楚很多。

先看原始 DeepSeekMath 的核心直觉。它把 PPO 中“用 value function 做 baseline”的部分，替换成“用同组样本的平均奖励做 baseline”。设同一个问题 $q$ 采样出 $G$ 个回答 $o_1,\dots,o_G$，对应奖励为 $r_1,\dots,r_G$。最常见的组内优势写法是：

\[A_i=\frac{r_i-\operatorname{mean}(r_{1:G})}{\operatorname{std}(r_{1:G})}\]

如果是 outcome supervision，那么一条回答中所有 token 共用这个终局归一化奖励；如果是 process supervision，则可以给每个推理步骤打分，再把未来步骤的归一化奖励累加到每个 token 上。DeepSeekMath 对 outcome supervision 和 process supervision 都给了明确说明。到这一步为止，重点是“优势来自组内相对奖励，而不是 value model”。

接下来，为了和 PPO 对照，很多讲解或实现会再引入 PPO-style 的概率比率。令

\[\begin{aligned} \rho_{i,t}(\theta) &= \frac{\pi_\theta(o_{i,t}\mid q,o_{i,<t})} {\pi_{\theta_{\text{old}}}(o_{i,t}\mid q,o_{i,<t})} \end{aligned}\]

则一种常见的 PPO-style surrogate 写法 可以写成：

\[\begin{aligned} \mathcal L_{\text{GRPO}}(\theta) &= -\mathbb E\Big[ \min\big( \rho_{i,t}A_i,\; \operatorname{clip}(\rho_{i,t},1-\epsilon,1+\epsilon)A_i \big) -\beta D_{\mathrm{KL}}(\pi_\theta\|\pi_{\text{ref}}) \Big] \end{aligned}\]

其中 $\beta$ 控制与参考策略的距离，$\epsilon$ 控制剪切范围。你会发现，它和 PPO 的结构非常像；区别在于 advantage 的来源：PPO 用 value/GAE，GRPO 用组内相对奖励。正因如此，你可以把 GRPO 记成“PPO 的 critic-free、group-relative 版本”。

最后要把“原始算法”和“工程变体”明确区分开。TRL 当前文档明确指出，组内标准差缩放可能引入 question-level difficulty bias；原始按回答长度做归一化的写法也被后续工作指出可能引入 length bias；因此当前工程实现会提供 scale_rewards="group" / "batch" / "none"、loss_type="grpo" / "dapo" / "dr_grpo"、mask_truncated_completions 等选项来缓解这些偏差。也就是说，网上你看到的 GRPO 公式之所以不完全一致，往往不是谁写错了，而是讨论的对象已经从“原始 DeepSeekMath 公式”切换成“后续工程修正版”。

伪代码

下面的伪代码表达的是“对每个 prompt 生成多个回答，再用组内相对优势更新”的核心思想。它与 DeepSeekMath 的算法框架和 TRL 的 GRPOTrainer 设计是一致的。

initialize policy πθ
initialize reference policy πref   # optional in some modern variants

for each iteration:
    batch_prompts = sample_prompts()

    groups = []
    for q in batch_prompts:
        completions = sample_G_responses(πθ_old, q, G)
        rewards = reward_fn(q, completions)   # rule-based or RM-based
        advantages = normalize_within_group(rewards)
        groups.append((q, completions, rewards, advantages))

    for epoch in range(num_iterations):
        for q, completions, rewards, advantages in groups:
            ratio = πθ(token | prefix) / πθ_old(token | prefix)
            loss = -mean(min(ratio * advantages,
                             clip(ratio, 1-eps, 1+eps) * advantages)
                         - beta * kl(πθ, πref))
            update(θ, loss)

简短 Python 示例片段

下面这段代码来自 TRL 的最小 GRPO 用法思路：定义数据集，定义奖励函数，把它交给 GRPOTrainer。这对初学者非常友好，因为它把“奖励函数”显式暴露出来了。

from datasets import load_dataset
from trl import GRPOTrainer
from trl.rewards import accuracy_reward

dataset = load_dataset("trl-lib/DeepMath-103K", split="train")

trainer = GRPOTrainer(
    model="Qwen/Qwen2-0.5B-Instruct",
    reward_funcs=accuracy_reward,
    train_dataset=dataset,
)

trainer.train()

实践要点与超参数建议

如果你是初学者，我最推荐把 GRPO 用在可验证任务上，而不是开放式写作上。原因很简单：当 reward 是“答案是否正确”“测试是否通过”“格式是否满足模板”时，组内相对比较非常自然，而且不容易被一个脆弱的神经奖励模型带偏。DeepSeek-R1 也正是在 reasoning 任务上强调使用 rule-based reward 与 format reward，并明确写出他们避免大规模使用神经奖励模型，是因为担心 reward hacking 和额外复杂度。

从 TRL 当前默认值看，GRPO 的入门配置相对克制：learning_rate 默认 1e-6，num_generations 默认 8，max_completion_length 默认 256，epsilon 默认 0.2，num_iterations 默认 1，beta 默认 0.0。文档还特别解释了：当 beta=0.0 时，参考模型根本不加载，这可以进一步降低显存与训练开销。与此同时，DeepSeekMath 在原始论文里给出过一个更“经典 RL”风格的配置：policy learning rate 取 1e-6，KL coefficient 取 0.04。因此，一个很实用的理解是：原始 GRPO 更像“带 KL 的 PPO 变体”，而许多后续实现开始尝试把 KL 弱化甚至去掉。

另一个必须注意的点是奖励归一化。TRL 文档直接提醒：按组标准差缩放虽然是原始直觉，但可能引入题目难度偏差；可选替代是 scale_rewards=False 或 scale_rewards="batch"。如果你在一个 batch 里混合了难度差异很大的题目，这个设置会明显影响训练信号。对刚开始实验的人，我建议先使用默认方式跑一次，再在一个小验证集上比较 group、batch、none 三种缩放策略，而不是一开始就盲信某一篇经验帖。

最后，GRPO 在长回答任务上要格外注意长度偏差与截断问题。TRL 当前已经把 mask_truncated_completions=True 明确标为一个有助于稳定性的好实践；同时也提供了针对长度偏差修正的 dr_grpo 等 loss 类型。你不必在第一次实验就把这些都用上，但必须知道：原始 GRPO 不是“公式一抄就稳”的算法，长 CoT 场景下很多偏差都是真问题。

典型实现与在线文档

GRPO 最核心的原始材料是 DeepSeekMath 论文；如果你想看它如何进入更大规模 reasoning 体系，则读 DeepSeek-R1。工程上，TRL 的 GRPOTrainer 是最适合初学者的公开实现接口，而 Hugging Face open-r1 则更接近“开放复现 DeepSeek-R1 路线”的研究/工程仓库。

中文资源里，优先看 Hugging Face 中文课程中实现 GRPO 的章节 和 OpenRLHF 中文文档。前者偏教学，后者偏工程系统。

常见问题与调试建议

GRPO 的第一类问题是“组内奖励几乎一样，学不到东西”。从公式直接可见，当同组样本的奖励全接近相同，归一化后的 advantage 就会接近 0，更新信号自然变弱。这通常不是算法坏了，而是你的 reward 太粗糙、组内采样多样性太低，或者任务本身对当前模型来说已经过于简单。

第二类问题是“算力一下子爆炸”。GRPO 每个 prompt 要采样 $G$ 个 completion，训练成本天然受 num_generations 影响；TRL 的 quick start 甚至直接写出，8 GPU 下一个示例训练大约需要 1 天。OpenRLHF 也强调，在 RLHF 系统里，生成阶段常是主要开销来源。对入门实验来说，你完全可以先把 num_generations、max_completion_length 和 batch size 压小，把奖励函数和日志跑通，再考虑扩展。

第三类问题是“你学到的是格式，而不是能力”。如果奖励函数只奖励 <think>...</think><answer>...</answer> 之类模板，模型很可能先学会的是“写得像被奖励的样子”，而不是更强的推理能力。因此 reward 设计最好至少同时覆盖正确性与格式性，这也是 DeepSeek-R1 中 accuracy reward 与 format reward 并用的原因。

小结与适用场景

如果你的任务是数学、代码、形式化推理、结构化输出，且有可靠 verifier，GRPO 非常值得学；如果你的任务是开放式闲聊、主观写作或通用帮助性，GRPO 也能用，但它的奖励设计会比 DPO 和经典 RM-PPO 更难做稳。从方法定位上看，GRPO 更接近“面向可验证推理任务的在线 RL 变体”。

方法比较与流程概览

比较表

方法 目标 是否需奖励模型 是否需价值模型 训练复杂度 优点 缺点 适用场景
RLHF 管线 用人类偏好把模型行为对齐到“更有帮助/更安全/更符合偏好”的方向 通常需要 取决于具体优化器；用 PPO 时通常需要 高 框架完整、可插入多种奖励、工业实践成熟 数据与工程链路长，调试成本高 通用大模型对齐、需要在线闭环时
PPO 稳定地在线优化策略，避免一步更新过大 在 LLM 对齐中通常需要奖励信号，常来自 RM 或规则 需要 高 通用、成熟、适合复杂奖励 actor/ref/reward/value 四件套很重 经典 RLHF、复杂环境、在线优化
DPO 直接从偏好对中优化策略 不需要显式 RM 不需要 中低 简洁、稳定、无需在线采样回路 强依赖高质量 preference pairs；$\beta$ 敏感 静态偏好数据、低门槛对齐
GRPO 用组内相对奖励做在线策略优化 可选：可用 RM，也可用规则/验证器 不需要 中高 去掉 critic，适合可验证推理任务 多候选采样成本高；公式和工程变体较多 数学、代码、结构化推理、规则奖励

表里的“是否需要奖励模型/价值模型”都要按语境理解。RLHF 是管线，不是单算法；如果 RLHF 用 PPO，则通常需要 RM 和 value model。DPO 不需要显式 RM，也不需要 value model。GRPO 在原始 DeepSeekMath 中保留了可选 KL 正则，但核心思想就是用组内相对奖励替代 critic；在 TRL 当前实现里，beta=0.0 甚至是默认值。

一张帮助记忆的流程对比图

flowchart TB
    subgraph A[经典 RLHF]
        A1[SFT 模型] --> A2[采样回答]
        A2 --> A3[偏好比较]
        A3 --> A4[奖励模型]
        A4 --> A5[PPO 更新策略]
    end

    subgraph B[DPO]
        B1[SFT 模型] --> B2[偏好比较数据]
        B2 --> B3[直接优化 chosen/rejected 概率边界]
    end

    subgraph C[GRPO]
        C1[SFT 或基座模型] --> C2[同一 prompt 采样 G 个回答]
        C2 --> C3[规则/奖励函数打分]
        C3 --> C4[组内相对优势]
        C4 --> C5[更新策略]
    end

这张图对应的记忆口诀可以很简单：RLHF 是全流程，PPO 是重武器，DPO 是轻武器，GRPO 是推理场景里常见的 critic-free 在线武器。这种归纳是基于原始论文与官方文档做的教学性总结，不是任何单一论文的原话。

参考链接

原始论文

  • Training language models to follow instructions with human feedback，InstructGPT 原始论文。
  • Learning to summarize from human feedback，OpenAI 在摘要任务上的经典 RLHF 论文。
  • Proximal Policy Optimization Algorithms，PPO 原始论文。
  • High-Dimensional Continuous Control Using Generalized Advantage Estimation，GAE 原始论文。
  • Direct Preference Optimization: Your Language Model is Secretly a Reward Model，DPO 原始论文。
  • DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models，GRPO 最核心的原始论文来源。
  • DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning，展示 GRPO 在 R1 系列中的核心地位与规则奖励实践。
  • The N+ Implementation Details of RLHF with PPO: A Case Study on TL;DR Summarization，公开复现 RLHF-with-PPO 的工程论文。
  • $\beta$-DPO: Direct Preference Optimization with Dynamic $\beta$，讨论 DPO 中 $\beta$ 敏感性的后续工作。

官方博客与官方文档

  • OpenAI: Aligning language models to follow instructions，InstructGPT 官方博客。
  • OpenAI: Learning to summarize with human feedback，总结任务 RLHF 官方介绍。
  • OpenAI Spinning Up，策略优化和 PPO 的权威入门材料。
  • Hugging Face TRL 总览，统一了解 PPO / DPO / GRPO / Reward Modeling 的入口。
  • TRL PPOTrainer 文档。
  • TRL DPOTrainer 文档。
  • TRL GRPOTrainer 文档。
  • TRL Reward Modeling 文档。
  • The N Implementation Details of RLHF with PPO，Hugging Face 对 OpenAI 早期 RLHF 代码的工程级复现总结。

官方或常用开源仓库

  • openai/following-instructions-human-feedback，OpenAI 的 InstructGPT 公开仓库，主要提供论文、模型卡和评测样本；完整训练代码未在该仓库中提供。
  • openai/lm-human-preferences，OpenAI 公开的早期“人类偏好微调语言模型”代码库。
  • openai/summarize-from-feedback，OpenAI 的摘要 RLHF 代码与数据资源。
  • openai/spinningup，OpenAI 的深度强化学习教学仓库。
  • eric-mitchell/direct-preference-optimization，DPO 参考实现。
  • huggingface/trl，当前最常见的 LLM 对齐训练工具库之一。
  • OpenRLHF/OpenRLHF，高性能 RLHF 工程框架。
  • deepseek-ai/DeepSeek-Math，DeepSeekMath 公开仓库。
  • huggingface/open-r1，Open-R1 开放复现项目。

中文高质量资源

  • Hugging Face 中文：从 RLHF 到 DPO。
  • Hugging Face 中文：DPOTrainer 文档。
  • Hugging Face 中文：使用 DPO 微调 Llama 2。
  • Hugging Face 中文课程：在 TRL 中实现 GRPO。
  • Spinning Up 中文版。
  • OpenRLHF 中文文档。

额外说明

  • 对于 InstructGPT 的完整训练代码仓库，本文检索到的是 OpenAI 的公开说明性仓库与相关 RLHF 参考代码；若你需要“完全复刻 InstructGPT 原训练流水线”的逐步脚本，公开材料中并未给出一套与论文一一对应的官方训练仓库，可视为未指定。
]]>JAX and TPU学习记录2026-03-26T00:00:00+00:002026-03-26T00:00:00+00:00https://wqh011128.github.io/blog/2026/03/26/JAX_and_TPU%E8%AE%B0%E5%BD%95JAX + TPU 两个重点问题

这次其实只需要抓住两个问题：

  1. 如何监控：我到底有没有真正用上 TPU，运行时状态怎么看？
  2. XLA_FLAGS 和 LIBTPU_INIT_ARGS 到底分别控制什么，会额外产出什么？

1. 如何监控 JAX + TPU

我现在更倾向于把“监控”分成三层看，因为很多混淆都来自把这三层混在一起。

1.1 第一层：机器上有没有 TPU 设备

这一层回答的是：

机器有没有挂上 TPU 芯片。

常用命令：

tpu-info
tpu-info --streaming --rate 2

如果 tpu-info 能看到型号、芯片数量、/dev/vfio/*，说明设备层面是存在的。

但这里要注意一个关键点：

  • 能看到 TPU 芯片，不等于能看到 TPU 利用率
  • HBM Usage、Duty cycle、TensorCore Utilization 出现 N/A，通常说明 runtime 指标没有正确暴露出来

最常见原因有这些：

  • 当前环境没有正确接上 libtpu
  • tpu-info 和当前环境不兼容
  • 当前没有真正跑 TPU workload
  • runtime metrics 本身还不可用

可以先做一个最小检查：

python - <<'PY'
try:
    import libtpu
    print("libtpu OK:", libtpu.__file__)
except Exception as e:
    print("libtpu import failed:", repr(e))
PY

如果这里连 libtpu 都导不进来，那么 tpu-info 读不到完整指标就不奇怪。

1.2 第二层：JAX 有没有成功连到 TPU backend

这一层回答的是：

JAX 这个进程，是否真的初始化了 TPU backend。

最小判断代码：

import jax

print("backend:", jax.default_backend())
print("devices:", jax.devices())

assert jax.devices()[0].platform == "tpu"

如果输出类似：

backend: tpu
devices: [TpuDevice(...)]

说明的事情只有一件：

JAX backend 已经接到 TPU 了。

但这还不等于“整段 Python 程序都在 TPU 上跑”。
真正上 TPU 的，只是 JAX/XLA 编译后的数组计算部分；数据加载、Python 循环、日志打印仍然主要在 host 侧。

1.3 第三层：具体这一步计算是否真的在 TPU 上执行

这一层回答的是：

不是环境认到了 TPU，而是这一步 step()、这次 matmul、这次训练迭代，是否真的在 TPU 上算了。

最实用的判断方式有三种。

方法 1：看输出数组落在哪个设备上

import jax
import jax.numpy as jnp

@jax.jit
def step(x):
    return x @ x + 1

x = jnp.ones((4096, 4096))
y = step(x)
y.block_until_ready()

print("device:", y.device)
print("sharding:", y.sharding)
print("addressable shard devices:", [s.device for s in y.addressable_shards])

如果 device 或 addressable_shards 对应的是 TpuDevice(...)，说明这一步结果确实落在 TPU 上。

方法 2：显式同步 block_until_ready()

这是判断 TPU 计算是否真的完成时最容易漏掉的一步。

JAX 默认是 asynchronous dispatch。也就是说：

  • Python 线程可能只是把任务提交给 TPU
  • 任务还在设备端排队或执行
  • 主线程已经继续往下跑了

所以像下面这种代码：

y = step(x)
print("done")

很多时候只能说明 dispatch 完成了，不能说明 TPU 已经算完了。

如果要做准确判断或准确计时，应该这样写：

import time
import jax
import jax.numpy as jnp

@jax.jit
def step(x):
    return x @ x

x = jnp.ones((4096, 4096))

t0 = time.time()
y = step(x)
t1 = time.time()

y.block_until_ready()
t2 = time.time()

print("dispatch time:", t1 - t0)
print("real execution time:", t2 - t0)

通常：

  • dispatch time 只是任务提交时间
  • real execution time 才更接近真实 TPU 执行时间

方法 3：需要铁证时，用 profiler

如果要看更硬的证据，直接做 profile。

import jax

jax.profiler.start_server(9999)

或者：

import jax

jax.profiler.start_trace("/tmp/profile-data")
# run workload
jax.profiler.stop_trace()

这时候关心的就不是“数组在哪”，而是 trace 里有没有真正的 TPU activity。

1.4 监控时最容易误判的三件事

误判 1：tpu-info 有芯片信息，就以为 TPU 一定在正常工作

不对。
这最多说明设备存在，不说明 runtime metrics 可用，更不说明 JAX 已经在跑 TPU 计算。

误判 2：程序很快打印 done，就以为 TPU 已经算完

不对。
这通常只是 asynchronous dispatch 的表现，必须配合 block_until_ready() 看。

误判 3：报 The TPU is already in use by process with pid ...，以为是 JAX 坏了

这类报错的核心含义其实很简单：

另一独立进程已经先初始化并占住了 TPU runtime。

常用检查命令：

ps -fp 22589
sudo lsof -w /dev/vfio/*
tpu-info

这通常是另一个脚本、另一个终端、tmux、notebook kernel 还没退出，不是 TPU “坏了”。

1.5 一个最小但可靠的监控顺序

如果我只想快速判断 TPU 状态，我会按这个顺序看：

  1. tpu-info 先确认机器层面有没有 TPU 设备。
  2. jax.default_backend() 和 jax.devices() 确认当前 JAX 进程有没有成功连到 TPU backend。
  3. 输出数组的 device / sharding / addressable_shards 确认具体计算结果是不是落在 TPU 上。
  4. block_until_ready() 确认这一步不是只完成 dispatch，而是真的执行完。
  5. profiler 需要最硬证据时再上。

2. HLO / LLO / compile / backend 到底是什么关系

如果只想抓主线，可以先记这一条：

trace -> lower -> HLO -> XLA 优化 -> backend lowering(LLO) -> executable

这一节其实就是在解释这条链路，以及它在项目里意味着什么。

2.1 四个词先分清

HLO 是什么

HLO 可以理解成 XLA 里的高层中间表示。
它更接近张量计算图，也更适合看：

  • 图优化
  • fusion
  • sharding
  • layout 传播

在项目里，下面两种写法都会碰到 HLO：

lowered_func.as_text("hlo")
lowered_func.compile().as_text()

区别是：

  • as_text("hlo") 更像是直接查看 lower 后的 HLO
  • compile().as_text() 则是在真正编译之后查看结果

compile() 是什么

compile() 不是“把内容打印出来”，而是：

  • 真正触发一次编译
  • 让 lowered computation 继续进入 XLA 和 backend 的编译链
  • 最终生成可执行对象

所以这句代码：

lowered_func.compile().as_text()

它的关键点不在 as_text()，而在 compile() 已经把后面的编译过程跑起来了。

backend 是什么

这里的 backend 可以理解为：

负责把编译器中间表示继续变成目标设备可执行结果的那一层。

在这个问题里，主要就是 TPU backend / libtpu。

LLO 是什么

LLO 可以粗略理解成更靠近后端和设备的一层表示。
它的位置比 HLO 更低，更接近最终 executable。

所以可以把它粗略记成：

HLO -> 更低层后端表示(LLO) -> executable

LLO 通常不是项目自己构造的数据结构，而是 TPU backend / libtpu 在编译过程中额外吐出来的调试产物。

2.2 这对当前项目意味着什么

当前项目里，如果只是：

lowered_func.as_text("hlo")

更偏向“看一份 HLO 文本”。

如果是：

lowered_func.compile().as_text()

那就不只是看文本了，而是会真的触发一次编译过程。
也正因为如此，编译链后面的很多东西都会被触发出来，比如：

  • 优化后的 HLO
  • XLA_FLAGS 对应的 HLO pass dump
  • LIBTPU_INIT_ARGS 对应的 LLO dump

所以项目里“导出优化后 HLO”这件事，本质上已经进入了真实编译，而不只是字符串导出。

2.3 为什么会同时看到 xla_dump 和 llo_dump

因为一次 compile() 本来就会经过多层编译链路。

如果设置：

XLA_FLAGS="--xla_dump_to=/tmp/xla_dump --xla_dump_hlo_pass_re=.*"
LIBTPU_INIT_ARGS="--xla_jf_dump_to=/tmp/dump_llo"

那么通常会同时出现两类 dump：

  • xla_dump 更偏 XLA / HLO pass，适合看编译器中高层是怎么改图的。
  • llo_dump 更偏 TPU backend / libtpu，适合看更靠近设备的一层 lowering。

它们不是重复关系，而是编译链上不同阶段的调试产物。

2.4 XLA_FLAGS 和 LIBTPU_INIT_ARGS 分别控制什么

如果只看用途，可以直接记成：

  • XLA_FLAGS：看 HLO / XLA pass
  • LIBTPU_INIT_ARGS：看 TPU backend / LLO lowering

更重要的一点是：

它们通常不是在改变最终 HLO 的语义，而是在让你看到更多编译过程中的中间状态。

所以同时打开时，通常会有三类结果：

  1. 项目自己主动导出的 HLO 文本
  2. /tmp/xla_dump 里的 HLO pass dump
  3. /tmp/dump_llo 里的 TPU/LLO dump

前者适合做稳定分析，后两者适合排查“图到底是在哪一层被改掉的”。

2.5 为什么不能在每个 layer 的 compile() 前动态改 LIBTPU_INIT_ARGS

直觉上很容易想到这种写法：

os.environ["LIBTPU_INIT_ARGS"] = "..."
lowered_func.compile()
os.environ.pop("LIBTPU_INIT_ARGS", None)

但这通常不可靠，原因很简单：

  • LIBTPU_INIT_ARGS 属于 backend / runtime 初始化参数
  • 它往往在 backend 初始化时就已经读取了
  • 不是每次 compile() 都重新读一次

所以程序跑到 save_hlo_and_data() 时，backend 很可能早就初始化完了。
这时再临时改环境变量，大概率不会影响当前这次 compile()。

因此它更像是：

  • 进程启动前配置
  • backend 初始化前配置

而不是 layer 级别的 Python 参数。

2.6 这次修改真正解决的问题

当前项目在导出优化后 HLO 时，会触发真实编译：

lowered_func.compile().as_text()

这次修改的目标不是改变编译行为本身，而是：

  • 继续保留全局 LIBTPU_INIT_ARGS="--xla_jf_dump_to=..."
  • 在每个 layer 的 compile() 前后观察全局 LLO dump 目录
  • 把本轮新增或变化的文件归档到当前 layer 目录

也就是把“全局 dump”尽量整理成“按 layer 可读”。

2.7 当前项目采用的最小方案

因为 LIBTPU_INIT_ARGS 只能全局生效，所以当前最稳妥的办法不是“让 libtpu 直接按 layer dump”，而是：

全局 dump，再由项目代码按每次 compile() 的时间窗口做归档。

具体流程是：

  1. compile 前读取全局 LLO dump 根目录
  2. 先做一次快照
  3. 执行当前 layer 的 compile()
  4. compile 后再做一次快照
  5. 找出本轮新增或变化的文件
  6. 复制到 output/<layer_name>/llo_dump

这个方案的优点是：

  • 不改 backend 初始化逻辑
  • 不需要多进程
  • 改动面小

它的限制也很明确：

  • 这是按“时间窗口”归因，不是按文件内容精确识别 layer
  • 如果 backend 异步写文件，可能有少量归因误差
  • 如果多个重复 block 共用同一个名字，仍然不容易区分实例

2.8 运行时怎么用

如果要让 LLO 归档生效，运行前需要先设置全局环境变量：

LIBTPU_INIT_ARGS="--xla_jf_dump_to=/tmp/dump_llo" \
uv run entrypoints.py \
  --model_type llama3 \
  --config-path configs/llama/llama3-8B-sq4k-bf16-L1.ini \
  --output_dir ./outputs/ \
  --export_after_optimize

如果还想同时看 HLO pass dump，再加：

XLA_FLAGS="--xla_dump_to=/tmp/xla_dump --xla_dump_hlo_pass_re=.*"

运行后通常会有两类目录：

  • 全局原始目录：/tmp/dump_llo
  • 项目归档目录：./outputs/<config_name>/<layer_name>/llo_dump

3. 最后压缩成几句话

  1. 监控 TPU 要分三层：设备是否存在、JAX 是否连上、具体计算是否真的执行。
  2. compile() 会真正触发编译，所以它不只是在“导出文本”，也会带出后端 dump。
  3. HLO 更偏 XLA 高层表示，LLO 更偏 TPU backend 更低层的 lowering。
  4. XLA_FLAGS 主要看 HLO/XLA pass，LIBTPU_INIT_ARGS 主要看 TPU backend/LLO。
  5. 当前项目不能让 libtpu 直接按 layer dump，所以采用的是“全局 dump + 每次 compile 后按 layer 归档”。
]]>FLUX.1梳理2026-03-25T00:00:00+00:002026-03-25T00:00:00+00:00https://wqh011128.github.io/blog/2026/03/25/FLUX.1FLUX.1 梳理

1. Flow Matching

Flow matching的核心在于——跳过加噪的过程。

常识建立：对于扩散模型的生成来说，假设$x$ 是原始分布，$x_1$ 是目标分布，$t$是timestep。则要求$x$ 能在$t$个step内将分布转移到$x_1$内，$x$ $\rightarrow$ $x_1$。

下面给出Flow matching的关键概念和直观理解。

1.1 ODE 与 SDE

ODE 的全称是 Ordinary Differential Equation，即常微分方程。
它用于描述一个变量随时间变化的确定性规律。一般形式可以写为： \(\frac{\mathrm{d}x}{\mathrm{d}t} = f(x,t)\)

其中：

  • $x$ 表示系统状态，在扩散模型中可以认为是某个时刻的数据分布；
  • $t$ 表示时间；
  • $f(x,t)$ 表示状态在当前时刻的变化率。

ODE 的特点是：在给定初始条件后，系统的演化轨迹是确定的，不包含随机扰动。

SDE 的全称是 Stochastic Differential Equation，即随机微分方程。
它用于描述既受到确定性规律影响、又受到随机噪声扰动的动态系统。一般形式可以写为： \(\mathrm{d}x = f(x,t)\,\mathrm{d}t + g(x,t)\,\mathrm{d}w\)

其中：

  • $f(x,t)\,\mathrm{d}t$ 是确定性的漂移项（drift term）；
  • $g(x,t)\,\mathrm{d}w$ 是随机扩散项（diffusion term）；
  • $w$ 表示 Brownian motion（布朗运动）。

SDE 的特点是：即使初始条件相同，由于存在随机项，系统轨迹通常也不是唯一的。

1.2 从ODE 出发，理解Flow Matching

再来看到ODE的公式： \(\frac{\mathrm{d}x}{\mathrm{d}t} = f(x,t)\) 那么，x对t的求导 $\frac{\mathrm{d}x}{\mathrm{d}t}$ 实际上表示状态变量 $x$ 关于时间 $t$ 的瞬时变化率，也就是$x$的变化方向，规定叫做Vector field。

在下图中，分布周围的蓝色箭头就是$f(x,t)$，可以直观理解为 “风吹着分布$x$从init状态$\rightarrow$target状态”。

规定$t\in [0,1]$，且将上面三张图分别定义init状态（t=0）、中间状态t时刻、和最终状态（t=1）。

则每一个step $t$时刻，数据的分布为$x_t$，vector field为$\frac{\mathrm{d}x_t}{\mathrm{d}t}$。

因此，可以得到一个初步结论：

在$t$个step内，从$x\rightarrow x_1$存在路径path，并且$f(x,t)$决定着$x_t$的走向。

[!Note]

所以，在扩散模型中，如果我们要将$x$生成到某个分布（例如猫和狗），意味着$x$和$x_1$已经确定，我们需要求解出合适的vector field，使得x移动的路径满足要求。

那么，$f(x,t)$该如何求解？先来看下面四张图：

图中的$\mu$就是$f$

  • （a）给定目标$x_1$情况下的Path条件分布，也即x的移动路径；
  • （b）不给定目标情况下，x移动路径的边缘分布，等于全部$x_1$下的期望；
  • （c）给定目标$x_1$情况下的vector field条件分布；
  • （d）不给定目标情况下，vector field的边缘分布；

对于（a），分布变化的过程等价于求解一条从$N(0,\sigma^2I)\rightarrow N(x_1,0)$的条件概率路径，

把正态分布参数化，定义成下面的形式： \(P_t(x|x_1):N(\alpha_{t}\cdot x_1,\beta_t^2\cdot I)\) 这样的好处是，可以将求解的过程转化成线性求解：

可以求出条件 vector field：

\[\begin{aligned} \frac{P_t(x\mid x_1)}{\mathrm{d}t} &= f_t(x\mid x_1) \\ &= \frac{\beta_t^\prime}{\beta_t}(x-\alpha_t) + \alpha_t^\prime \end{aligned}\]

边缘分布等于条件分布的期望，先铺垫以下概念：

  1. Divergence：在当前t step x的位置，流出的量减去流入的量（直观理解为当前位置，可以出去的方向数 减去 可以进来的方向数）
\[\operatorname{div}(v_t)(x)=\sum_{i=1}^{d}\frac{\partial}{\partial x_i}v_t(x)\]
  1. Continuity Equation：质量守恒。
\[\frac{\mathrm{d}}{\mathrm{d}t}p_t(x)+\operatorname{div}\bigl(p_t\cdot f_t\bigr)(x)=0\]

直观理解：

  • 当这个位置的散度（流出的量-流入的量）为0时，也就是$\operatorname{div}\bigl(p_t\cdot f_t\bigr)(x)$=0时，物体的质量随时间不会发生改变。

  • 如果散度大于0，那么质量会随时间减小，$\frac{\mathrm{d}}{\mathrm{d}t}p_t(x)$<0。
  • 如果散度小于0，那么质量会随时间变大，$\frac{\mathrm{d}}{\mathrm{d}t}p_t(x)$>0。

Proof:

\[\begin{aligned} \frac{\mathrm{d}}{\mathrm{d}t}p_t(x) &= \frac{\mathrm{d}}{\mathrm{d}t} \int p_t(x\mid x_1)\,p(x_1)\,\mathrm{d}x_1 \\ &= \int \frac{\mathrm{d}}{\mathrm{d}t}p_t(x\mid x_1)\,p(x_1)\,\mathrm{d}x_1 \\ &= \int \Bigl[ -\operatorname{div}\bigl(p_t(x\mid x_1)\,f_t(x\mid x_1)\bigr) \Bigr] p(x_1)\,\mathrm{d}x_1 \\ &= -\operatorname{div} \int p_t(x\mid x_1)\,f_t(x\mid x_1)\,p(x_1)\,\mathrm{d}x_1 \\ &= -\operatorname{div} \left[ p_t(x) \int f_t(x\mid x_1)\, \frac{p_t(x\mid x_1)\,p(x_1)}{p_t(x)} \,\mathrm{d}x_1 \right] \\ &= -\operatorname{div}\bigl(p_t\cdot f_t\bigr)(x) \end{aligned}\]

Therefore,

\[\begin{aligned} f_t(x) &= \int f_t(x\mid x_1)\, \frac{p_t(x\mid x_1)\,p(x_1)}{p_t(x)} \,\mathrm{d}x_1 \end{aligned}\]

1.3 How to train a Flow-Matching Model

在上一部分中，我们已经得到边缘向量场的表达式：

\[\begin{aligned} f_t(x) &= \int f_t(x\mid x_1)\, \frac{p_t(x\mid x_1)p(x_1)}{p_t(x)} \,\mathrm{d}x_1 \\ &= \mathbb{E}_{x_1\sim p(x_1\mid x,t)}\bigl[f_t(x\mid x_1)\bigr]. \end{aligned}\]

这说明，边缘向量场 $f_t(x)$ 本质上是条件向量场 $f_t(x\mid x_1)$ 关于后验分布的条件期望。

我们的目标是训练一个参数化模型 $f_t^\theta(x)$，使其能够逼近真实的边缘向量场 $f_t(x)$。
一个自然的训练目标是最小化如下损失： \(\begin{aligned} \mathcal{L}_{\mathrm{fm}}(\theta) &= \mathbb{E}_{t,x_1,x} \left[ \left\|f_t^\theta(x)-f_t(x)\right\|^2 \right]. \end{aligned}\)

但是这个目标通常是 intractable 的，因为 $f_t(x)$ 本身往往无法直接计算。

为了解决上述问题，考虑定义条件 Flow Matching 损失： \(\begin{aligned} \mathcal{L}_{\mathrm{cfm}}(\theta) &= \mathbb{E}_{t,x_1,x} \left[ \left\|f_t^\theta(x)-f_t(x\mid x_1)\right\|^2 \right]. \end{aligned}\)

下面说明两个损失函数 L_fm 和 L_cfm 仅相差一个与参数 theta 无关的常数，因此二者具有相同的最优解。

1. 展开 L_fm

\[\begin{aligned} \mathcal{L}_{\mathrm{fm}}(\theta) &= \mathbb{E}_{t,x_1,x} \left[ \|f_t^\theta(x)-f_t(x)\|^2 \right] \\ &= \mathbb{E}_{t,x_1,x} \left[ \|f_t^\theta(x)\|^2 -2f_t^\theta(x)^\top f_t(x) +\|f_t(x)\|^2 \right]. \end{aligned}\]

2. 展开 L_cfm

\[\begin{aligned} \mathcal{L}_{\mathrm{cfm}}(\theta) &= \mathbb{E}_{t,x_1,x} \left[ \|f_t^\theta(x)-f_t(x\mid x_1)\|^2 \right] \\ &= \mathbb{E}_{t,x_1,x} \left[ \|f_t^\theta(x)\|^2 -2\,f_t^\theta(x)^\top f_t(x\mid x_1) +\|f_t(x\mid x_1)\|^2 \right]. \end{aligned}\]

3. 利用边缘向量场的定义

由上一节结论，

\[\begin{aligned} f_t(x) &= \int p(x_1\mid x,t)\,f_t(x\mid x_1)\,\mathrm{d}x_1 \\ &= \mathbb{E}_{x_1\sim p(x_1\mid x,t)}\bigl[f_t(x\mid x_1)\bigr]. \end{aligned}\]

因此，对于固定的 $x,t$，有

\[\mathbb{E}_{x_1\sim p(x_1\mid x,t)} \left[ f_t(x)-f_t(x\mid x_1) \right] =0.\]

进一步可得

\[\mathbb{E}_{t,x} \left[ f_t^\theta(x)^\top \mathbb{E}_{x_1\sim p(x_1\mid x,t)} \bigl[f_t(x)-f_t(x\mid x_1)\bigr] \right] =0.\]

也就是

\[\begin{aligned} \mathbb{E}_{t,x_1,x} \left[ f_t^\theta(x)^\top f_t(x) \right] &= \mathbb{E}_{t,x_1,x} \left[ f_t^\theta(x)^\top f_t(x\mid x_1) \right]. \end{aligned}\]

4. 两个损失之差

于是， \(\begin{aligned} \mathcal{L}_{\mathrm{fm}}(\theta)-\mathcal{L}_{\mathrm{cfm}}(\theta) &= \mathbb{E}_{t,x_1,x} \left[ \|f_t(x)\|^2-\|f_t(x\mid x_1)\|^2 \right]. \end{aligned}\)

右边不含参数 $\theta$，因此它只是一个常数。于是有

\[\begin{aligned} \mathcal{L}_{\mathrm{fm}}(\theta) &= \mathcal{L}_{\mathrm{cfm}}(\theta)+\text{constant}. \end{aligned}\]

因此，

\[\begin{aligned} \arg\min_\theta \mathcal{L}_{\mathrm{fm}}(\theta) &= \arg\min_\theta \mathcal{L}_{\mathrm{cfm}}(\theta). \end{aligned}\]

所以，虽然我们原本想训练的是边缘向量场 $u_t(x)$，但实际中可以转而优化如下可计算的目标：

\[\begin{aligned} \mathcal{L}(\theta) &= \mathbb{E}_{t,x_1,x} \left[ \left\|f_t^\theta(x)-f_t(x\mid x_1)\right\|^2 \right]. \end{aligned}\]

这就是 Flow Matching / Conditional Flow Matching 的核心训练形式。

它的含义是：

  • 我们真正想拟合的是边缘向量场 $f_t(x)$；
  • 但由于 $f_t(x)$ 不易直接计算，于是改为拟合条件向量场 $f_t(x\mid x_1)$；
  • 由于二者对应的损失只差一个与参数无关的常数，因此训练条件目标就等价于训练边缘目标。

1.4 Pseudo codes of training and inference

线性情况下， $\frac{P_t(x|x_1)}{\mathrm{d}t}=f_t(x|x_1)=\frac{\beta_t^\prime}{\beta_t}(x-\alpha_t) + \alpha_t^\prime$

Training

1. 从数据集中采样一个图像 $x_1$。

2. 随机采样一个 timestep $t\in[0,1]$，以及一个 noise $\sigma\in N(0,I)$。

3. 给定 $x_1$ 下的条件路径分布为

\[p_t(x\mid x_1)=t x_1 + (1-t)\sigma\]

它服从分布

\[N(\alpha_t\cdot x_1,\beta_t^2\cdot I),\]

所以 $\alpha_t=t,\beta_t=1-t$。

4. 因此

\[f_t(x\mid x_1)=x_1-\sigma\]

或者你直接对 $t$ 求导，也是这个结果，不用管服从什么分布。

5. 所以损失函数为

\[\mathcal{L}(\theta)=\left\|f_t^\theta(x)-(x_1-\sigma)\right\|^2.\]

6. BP优化，最终模型能够拟合vector field。可以注意到不论 $t$ 采样多少，监督信号都是不变的。

Inference

1. 初始化 $t=0$，step size 为 $h=\frac{1}{n}$，其中 $n$ 为离散步数。

2. 令 $x_0$ 为初始状态。

3. 对 $i=1,\dots,n-1$ 循环执行：

\[x_{t+h}=x_t+f_t^\theta(x_t)\cdot h\] \[t\leftarrow t+h\]

4. 返回 $x_1$。

2. DiT

主要从这张图看一下时间步这种条件是如何影响模型的生成的，也就是$f_t^\theta(x_t)$怎么接受$x$和$t$。

在 DiT 中，若采用 adaLN-Zero 机制，则模型虽然形式上可写为 $f_t^\theta(x_t)$， 但更准确地说，它实际还依赖于时间步条件以及其他外部条件，因此应理解为 $f_\theta(x_t, t, c)$。

其中：

  • $x_t$ 表示时刻 $t$ 的带噪 latent；
  • $t$ 表示时间步或噪声强度；
  • $c$ 表示附加条件，例如类别标签或文本条件。
  1. $x_t$ 如何进入模型

输入的带噪 latent $x_t$ 先经过 patchify，被划分为一系列 token，作为 Transformer 主干网络的输入。
因此，$x_t$ 提供的是当前样本的空间内容信息，是模型生成更新的基础。

  1. $t$ 和条件 $c$ 如何进入模型

在 adaLN-Zero 中，时间步 $t$ 以及其他条件 $c$ 不直接作为 token 与输入序列拼接，也不通过 cross-attention 进入主干，而是先经过 embedding 和 MLP，生成一组用于调制 Transformer block 的参数。

设条件向量记为

\[e = \phi(t, c),\]

其中 $\phi(\cdot)$ 表示由时间步嵌入、条件嵌入及后续 MLP 构成的映射。
随后，模型根据该条件向量生成每个 block 所需的调制参数，例如：

\[(\alpha_1,\beta_1,\gamma_1,\alpha_2,\beta_2,\gamma_2)=\mathrm{MLP}(e).\]

这些参数分别作用于 block 中的 attention 分支和 feedforward 分支。

  1. adaLN-Zero 的具体作用方式

对于输入特征 $h$，普通 LayerNorm 可写为

\[\mathrm{LN}(h).\]

在 adaLN 中，LayerNorm 的输出会被条件相关参数进一步调制：

\[\mathrm{adaLN}(h,e)=\gamma(e)\odot \mathrm{LN}(h)+\beta(e),\]

其中：

  • $\gamma(e)$ 表示缩放参数（scale）；
  • $\beta(e)$ 表示平移参数（shift）；
  • $\odot$ 表示逐元素乘法。

此外，残差分支还会通过门控参数 $\alpha(e)$ 进行控制，因此对应的更新形式可以写为

\[h+\alpha(e)\cdot F\bigl(\mathrm{adaLN}(h,e)\bigr),\]

其中 $F(\cdot)$ 表示 self-attention 或 pointwise feedforward 等子模块。

  1. Zero 的含义

adaLN-Zero 中的 “Zero” 表示：在初始化时，调制相关的残差门控通常被初始化为接近零，使得网络初始状态更接近恒等映射。
这样做的作用是：

  • 降低训练初期条件注入对主干特征的剧烈扰动；
  • 提高深层 Transformer 训练的稳定性；
  • 使模型逐步学会如何利用时间步和条件信息调制生成过程。
  1. 对生成过程的理解

因此，在 adaLN-Zero 中：

  • $x_t$ 决定当前 noisy latent 的内容；
  • $t$ 告诉模型当前处于去噪过程的哪个阶段；
  • 条件 $c$（例如 text）提供目标语义信息。

但这些条件并不是直接替代输入特征，而是通过生成 scale、shift 和 gate 等参数，逐层调制 Transformer block 对 $x_t$ 的处理方式。
因此，$f_t^\theta(x_t)$ 的本质是一个在条件控制下对输入 latent 进行变换的函数，其条件依赖性由 adaLN-Zero 注入到每一层的归一化与残差更新过程中。

3. FLUX.1 Model

前面已经把 Flow Matching 和 DiT 中条件如何注入 两件事讲清楚了。

接下来就要回答真正和 FLUX.1 本身有关的问题：

  • 输入文本是怎么进入模型的？
  • 图像为什么不是直接在像素空间生成，而是在 latent 空间生成？
  • FLUX.1 和普通 DiT 到底有什么结构差别？
  • 模型每一步预测的到底是什么量？

如果把这些问题串起来，FLUX.1 的整体逻辑就会变得非常清楚。

3.1 总体架构总览

先给出一句最核心的概括：

FLUX.1 本质上是一个 latent rectified-flow transformer。

这句话可以拆成三层含义：

  1. latent

    模型不是直接在原始像素空间上生成图像，而是先在压缩后的 latent 空间中做生成，最后再 decode 回图像。

  2. rectified flow

    模型学习的是从噪声分布流向目标图像分布的 vector field，因此训练目标更接近前面讲的 Flow Matching，而不是传统 DDPM 里一步步预测噪声的写法。

  3. transformer

    主干网络不再是 UNet，而是将图像 latent 划分成 token，再和文本 token 一起送入 Transformer 中建模。

从 diffusers 的参考实现来看，一个完整的 FLUX pipeline 主要包括下面几个部分：

  • VAE：负责图像和 latent 之间的编码/解码；
  • text_encoder：CLIPTextModel，可以理解为提供更偏全局的文本语义；
  • text_encoder_2：T5EncoderModel，可以理解为提供更细粒度、更长序列的文本条件；
  • transformer：真正负责预测 latent 更新方向的 FluxTransformer2DModel；
  • scheduler：FlowMatchEulerDiscreteScheduler，负责按照 flow matching 的离散更新方式推进采样过程。

因此，FLUX.1 并不是“只换了个 loss 的 DiT”。

更准确地说，它是：

  1. 在 latent 空间中进行生成；
  2. 用 文本-图像联合 Transformer 作为主干；
  3. 用 flow / rectified flow 的目标来学习生成轨迹。

3.2 输入表示：text、latent 与 token

在真正进入 Transformer 之前，FLUX.1 需要先把“文本”和“图像”都表示成 token 形式。

3.2.1 图像侧：先变成 latent，再变成 image tokens

图像侧的流程可以理解为：

image -> VAE encoder -> latent grid -> image tokens

为什么不直接在像素空间生成？

原因很简单：像素空间太大，直接建模的计算量和显存压力都非常高。
因此和 Stable Diffusion 系列一样，FLUX.1 也是先在更低维的 latent 空间中完成主要生成过程，再由 VAE decoder 还原出最终图像。

在 diffusers 的 FluxTransformer2DModel 中，图像输入 hidden_states 的形状是：

\[(\text{batch},\ \text{image\_sequence\_length},\ \text{in\_channels})\]

并且参考实现中的 in_channels=64。
这说明进入 Transformer 时，图像已经不再是二维 feature map，而是被整理成一串 image tokens，每个 token 对应一个 latent 网格位置上的特征。

所以可以把 FLUX.1 的图像输入理解为：

  • 空间结构仍然保留在 token 的排列关系中；
  • 但网络实际处理的是一串序列，而不是传统卷积网络里的 feature map。

3.2.2 文本侧：不是一条文本通路，而是两条

FLUX pipeline 的另一个关键点在于：它并不是只用一个 text encoder。

在 diffusers 参考实现中，它同时使用：

  • CLIPTextModel
  • T5EncoderModel

可以把这两条文本通路粗略理解为：

  • CLIP 更偏向于提供一个全局、压缩后的语义摘要；
  • T5 更偏向于提供细粒度的文本序列信息，保留 prompt 中更多 token 级别的内容。

这两部分信息在后续主干中扮演的角色并不完全一样：

  • 一部分作为 序列条件 参与 joint attention；
  • 一部分作为 pooled condition，进入时间步/条件调制通路，用于控制 block 的行为。

这也是为什么 FLUX.1 的 prompt adherence 很强。
因为它并不是简单地把一句文本编码成一个向量，而是同时保留了：

  • 全局语义；
  • 长文本细节；
  • token 级别的结构信息。

3.3 FLUX.1 主干：先双流，再单流

如果只看名字，很多人会觉得 FLUX.1 就是一个普通 DiT。
但从 diffusers 中 FluxTransformer2DModel 的结构参数来看，它其实有一个非常鲜明的组织方式：

  • num_layers = 19
  • num_single_layers = 38

文档中把前者称作 dual stream DiT blocks，后者称作 single stream DiT blocks。

这件事非常重要，因为它揭示了 FLUX.1 的核心结构逻辑。

3.3.1 dual stream：先分别处理 text 和 image

所谓 dual stream，可以把它理解成：

  • 文本 token 有自己的流；
  • 图像 token 也有自己的流；
  • 但二者并不是完全隔离，而是在 attention 中交换信息。

为什么要这样设计？

因为 text token 和 image token 的统计特性并不一样：

  • text token 更偏离散语义；
  • image token 更偏连续空间特征。

如果一上来就把两者完全混在一起做统一建模，模型需要在很浅层就同时解决：

  • 文本语义理解；
  • 图像结构建模；
  • 跨模态对齐；

这会让优化变得更困难。

因此 FLUX.1 的策略是：

  1. 先保留 text/image 各自的表示流；
  2. 在较浅层完成跨模态对齐；
  3. 等两种 token 已经“知道彼此是谁”之后，再进入更深的统一建模。

3.3.2 single stream：后面再做深度融合

在 dual stream blocks 之后，FLUX.1 还有更多的 single stream blocks。

这可以理解为：

  • 前面解决“对齐”问题；
  • 后面解决“联合生成”问题。

一旦 text token 和 image token 已经在浅层对齐完成，后面的单流建模就可以更专注于：

  • 细化空间结构；
  • 强化 prompt 对局部内容的控制；
  • 输出更准确的 vector field。

因此，FLUX.1 的主干不是“从头到尾都一个结构”，而是分成两个阶段：

  1. dual stream 阶段：做跨模态对齐；
  2. single stream 阶段：做联合生成与精细更新。

这也是它和“普通一串统一 Transformer block”的一个重要差别。

3.4 FLUX.1 预测的到底是什么

这一点必须和传统 diffusion 区分开。

在很多 DDPM / latent diffusion 模型中，网络常见的预测目标是：

  • 噪声 $\epsilon$；
  • 或者某种等价的重参数化目标，例如 $v$。

但在 FLUX.1 里，更自然的理解方式是：

模型预测的是当前 latent 状态应该往哪个方向走，也就是 vector field / velocity。

因此，输入当前时刻的 latent $x_t$ 后，模型输出可以理解为：

\[f_t^\theta(x_t,\text{text})\]

它描述的是：在当前时间步，latent 应该沿哪个方向更新，才能逐渐流向目标图像分布。

所以推理时的核心更新形式就是：

\[x_{t+h}=x_t+h\cdot f_t^\theta(x_t,\text{text})\]

这和前面讲的 ODE 完全对应。

也就是说，FLUX.1 的生成过程更像是：

  • 在 latent 空间里定义一条从 noise 到 image 的轨迹；
  • Transformer 在每个时刻给出局部方向；
  • scheduler 按照这个方向做离散积分。

从这个角度看，FLUX.1 和“逐步去噪”的直觉并不完全一样。
它更像是在学习一张“速度场地图”，然后沿着这张地图把噪声流到图像。

3.5 条件信息到底怎么进入 FLUX.1

前面在 DiT 一节已经介绍过 adaLN-Zero 的基本思想。
到了 FLUX.1 这里，可以把条件注入方式理解得更具体一些。

对于 FLUX.1 来说，至少有三类条件会参与计算：

  1. 当前 latent 状态 $x_t$；
  2. 时间步 timestep；
  3. 文本条件（包括序列级文本表示和 pooled 文本表示）。

从 FluxTransformer2DModel.forward 的输入可以看到，主干会接收：

  • hidden_states
  • encoder_hidden_states
  • pooled_projections
  • timestep
  • 以及和位置编码有关的 img_ids、txt_ids

这说明 FLUX.1 的条件注入并不是“只有一条 text embedding”这么简单，而是至少分成了两条路径：

3.5.1 一条进入 attention

encoder_hidden_states 可以理解为文本 token 序列对应的条件表示。
这部分信息参与 joint attention，使图像 token 在更新时能够感知 prompt 的细粒度语义。

3.5.2 一条进入 block 调制

pooled_projections 和 timestep 则更像全局条件。
它们不会只告诉模型“要画什么”，还会告诉模型：

  • 当前处于采样轨迹的哪个阶段；
  • 当前 block 应该更偏向语义布局，还是更偏向细节修正。

因此，如果把 FLUX.1 看成一个受条件控制的动力系统，那么：

  • encoder_hidden_states 更像局部语义约束；
  • pooled_projections 更像全局控制信号；
  • timestep 更像系统当前所处的时间位置。

这三者一起决定了最终的 vector field。

3.6 dev 和 schnell 的差别

FLUX.1 常见的开源版本里，最常被拿来比较的就是：

  • FLUX.1-dev
  • FLUX.1-schnell

二者都基于 FLUX.1 主干，但推理习惯和蒸馏方式不同。

3.6.1 dev

根据 diffusers 文档，dev 是 guidance-distilled 版本。

它的典型特点是：

  • 生成质量更稳；
  • 常见设置下大约需要 50 steps 才能比较充分地发挥效果；
  • 没有 max_sequence_length=256 这类限制；
  • 推理时通常会配合非零 guidance_scale 使用。

因此，dev 更像是：

  • 速度不是第一优先级；
  • 更注重质量、细节和 prompt 跟随能力。

3.6.2 schnell

根据 diffusers 文档，schnell 是 timestep-distilled 版本。

它的特点更激进：

  • 通常 1 到 4 步 就能出图；
  • guidance_scale 需要设为 0；
  • max_sequence_length 不能超过 256。

因此，schnell 的本质是：
用额外蒸馏把原本多步的采样轨迹压缩到极少数步骤里，从而换取极快的生成速度。

可以简单理解为：

  • dev：更偏质量；
  • schnell：更偏速度。

3.7 FLUX.1 和 SDXL / 普通 DiT 的差别

如果要快速定位 FLUX.1 在整个生成模型谱系里的位置，可以这样看：

3.7.1 和 SDXL 的差别

SDXL 的主干仍然是 UNet 思路：

  • 图像特征在多尺度卷积结构里流动；
  • 条件通常通过 cross-attention 等方式注入。

而 FLUX.1 的主干已经变成：

  • 图像先转成 token；
  • 文本也转成 token；
  • 再用 Transformer 做统一建模。

所以二者最大的差别不是“loss 不一样”，而是：

  • SDXL 更像卷积式的 latent diffusion；
  • FLUX.1 更像 token-based 的 latent flow transformer。

3.7.2 和普通 DiT 的差别

普通 DiT 的核心思想是：
用 Transformer 代替 UNet 来处理图像 token。

而 FLUX.1 在此基础上又往前走了一步：

  • 不只是“图像 token + 条件”；
  • 而是显式地做 文本流 / 图像流的分阶段融合；
  • 再配合 rectified flow / flow matching 目标来学习生成轨迹。

所以从直观上说：

  • DiT 解决的是“能不能用 Transformer 生成图像”；
  • FLUX.1 更进一步解决的是“怎么让文本条件和图像 latent 在 Transformer 里更高效地协同工作”。

3.8 一个完整的 forward pass

最后，把前面的所有内容串起来，看一遍 FLUX.1 的完整前向流程。

假设输入是一句 prompt，例如：

"a tiny astronaut hatching from an egg on the moon"

则一次典型生成可以理解为下面几步：

3.8.1 文本编码

prompt 先分别送入文本编码器，得到两类文本条件：

  • 序列级文本 token 表示；
  • pooled 的全局文本表示。

3.8.2 初始化 latent

在图像侧，先采样一个初始噪声 latent，记作 $x_0$。
此时它还不对应任何真实图像，只是 latent 空间中的随机起点。

3.8.3 整理成 image tokens

这个 latent 会被整理成 image token 序列，作为 Transformer 的 hidden_states 输入。

3.8.4 注入时间步和条件

当前的 timestep、pooled 文本条件、文本 token 条件一起进入主干网络：

  • 一部分调制 block；
  • 一部分进入 attention；
  • 一部分决定 text/image 的对齐关系与融合方式。

3.8.5 Transformer 预测 vector field

经过 dual stream blocks 和 single stream blocks 后，模型输出当前时刻的更新方向，也就是：

\[f_t^\theta(x_t,\text{text})\]

3.8.6 scheduler 更新 latent

随后，FlowMatchEulerDiscreteScheduler 根据这个方向，把 latent 从当前时刻推进到下一个时刻：

\[x_{t+h}=x_t+h\cdot f_t^\theta(x_t,\text{text})\]

重复多次之后，latent 会逐渐从随机噪声流向目标图像分布。

3.8.7 decode 成最终图像

当采样结束后，得到最终 latent $x_1$，再送入 VAE decoder，输出最终图像。

3.9 从工程角度如何理解 FLUX.1

如果从工程实现的角度看，FLUX.1 之所以强，主要是因为它同时把三件事做到了：

  1. 文本理解强

    双文本编码器 + 联合 Transformer，使 prompt 中的复杂语义能够更充分地注入图像生成过程。

  2. 主干表达能力强

    12B 级别的 Transformer 主干，加上先 dual-stream 再 single-stream 的结构，使它在跨模态对齐和图像细节生成上都有很高容量。

  3. 采样路径更直接

    flow matching / rectified flow 的目标让模型更自然地学习“从噪声走向图像”的方向场，因此在推理时可以用更直接的轨迹更新 latent。

当然，它的代价也很明显：

  • 模型非常大；
  • 显存压力高；
  • 推理成本高；
  • consumer GPU 上部署并不轻松。

所以可以把 FLUX.1 看成一个典型的“拿更大算力和更复杂结构，换更强 prompt adherence 与更高图像质量”的模型。

到这里再回头看，你会发现 FLUX.1 的逻辑其实很统一：

  • 理论层面：用 Flow Matching 学 vector field；
  • 网络层面：用 Transformer 建模 text/image token；
  • 条件层面：用时间步和文本条件调制主干；
  • 推理层面：在 latent 空间中沿着 learned flow 逐步更新。

这四件事合在一起，才构成了完整的 FLUX.1。

]]>Git 常见问题与操作笔记2026-03-11T00:00:00+00:002026-03-11T00:00:00+00:00https://wqh011128.github.io/blog/2026/03/11/GIT这篇笔记按“遇到的问题”整理 Git 日常操作。命令默认在项目根目录执行；如果涉及已经推送到远端的历史改写，优先用 --force-with-lease，不要直接用 --force。

GitHub recovery codes 应该写进笔记里吗？

不要。GitHub recovery codes 等同于备用登录凭证，不应该放在仓库、博客、截图或公开笔记里。

如果 recovery codes 曾经被推送到公开仓库或 GitHub Pages，建议在 GitHub 的 2FA 设置里重新生成 recovery codes，让旧 codes 失效。

怎么把 remote 的 HTTPS 地址换成 SSH？

查看当前 remote：

git remote -v

把 origin 从 HTTPS 改成 SSH：

git remote set-url origin git@github.com:<yourname>/<repo>.git

再次确认：

git remote -v

Git 的全局身份、提交模板和 SSH key 怎么配置？

身份信息：

git config --global user.name "wqh011128"
git config --global user.email "wqh011128@163.com"

提交信息模板：

cat <<'EOF' > ~/.gitmessage.txt
<type>[<scope>]: <short-summary>

Problem:
<description of the problem being solved>

Solution:
<description of the solution implemented>

Test:
<description of how the change was tested>

JIRA: ISSUE-<Number>
EOF

git config --global commit.template ~/.gitmessage.txt

Windows 上如果经常遇到 CRLF/LF 问题，可以按项目规范决定是否设置：

git config --global core.autocrlf true

生成 SSH key 并查看公钥：

ssh-keygen -t ed25519 -C "wqh011128@163.com"
cat ~/.ssh/id_ed25519.pub

如果环境不支持 ed25519，再考虑 RSA：

ssh-keygen -t rsa -b 2048 -C "wqh011128@163.com"
cat ~/.ssh/id_rsa.pub

git fetch、git pull 和 git pull --rebase 有什么区别？

git fetch 只更新远端分支信息，不会改动当前工作区和当前分支：

git fetch origin

git pull 等价于先 fetch，再把远端内容整合到当前分支。默认通常是 merge：

git pull origin main

这条命令的意思是：把远端 origin/main 整合到你当前 checkout 的分支里。它不是“切到 main 再拉取”，当前分支是谁就影响谁。

如果想让历史更线性，常用 rebase：

git pull --rebase origin main

概念上可以理解成：

git fetch origin main
git rebase origin/main

git pull --rebase origin dev/vl 会改哪个分支？

会改你当前 checkout 的分支。

例如你当前在 feat 分支，执行：

git pull --rebase origin dev/vl

含义是：

  1. 从 origin 拉取 dev/vl 的最新提交。
  2. 把当前分支 feat 上你自己的提交 rebase 到 origin/dev/vl 之后。

所以它不是“把 dev/vl 拉到 dev/vl”，而是“用 origin/dev/vl 作为当前分支的新基底”。

例子：

origin/dev/vl: A-B-C
local feat:    x-y

执行后大致变成：

A-B-C-x'-y'

怎么复制一个 branch 继续开发？

如果想基于最新 main 开一个新分支：

git switch main
git pull --rebase origin main
git switch -c feat/example

老版本 Git 可以用：

git checkout -b feat/example

怎么查看、添加和删除 remote？

查看远程连接：

git remote -v

添加远程连接：

git remote add origin <remote-url>

删除某个远程连接：

git remote remove origin

修改远程连接：

git remote set-url origin <remote-url>

怎么从本地项目创建 GitHub 仓库并推上去？

如果本地还没有 Git 历史：

cd /path/to/your-project
git init
git add .
git commit -m "Initial commit"

如果本地已经有提交历史，从绑定远端开始即可。

在 GitHub 上创建仓库时，建议创建空仓库：

  • 填仓库名。
  • 先选 Private 或 Public 都可以，按需求决定。
  • 不要勾选 README、.gitignore、license，避免和本地历史冲突。

使用 HTTPS：

git remote add origin https://github.com/<yourname>/<repo>.git
git branch -M main
git push -u origin main

使用 SSH：

git remote add origin git@github.com:<yourname>/<repo>.git
git branch -M main
git push -u origin main

常见坑：

  • HTTPS push 被拒绝或要求登录：GitHub 通常需要 Personal Access Token，不是账号密码。
  • 远端已有 README 导致冲突：可以重建空仓库，或先 git pull --rebase 解决冲突后再 push。
  • 默认分支不是 main：用 git branch -M main 统一。

怎么创建新分支并推送到远端？

先确认当前状态：

git status
git branch --show-current

基于当前分支创建并切到新分支：

git switch -c feat/eau-opt

提交修改：

git add -A
pre-commit run --all-files
git commit -m "Describe your change"

推送新分支：

git push -u origin feat/eau-opt

-u 会建立 upstream。以后在这个分支上直接 git push 或 git pull 就行。

当前 feat 分支怎么更新到最新 origin/main？

推荐把当前 feature 分支 rebase 到最新 origin/main 上：

git status
git stash push -u -m "wip before rebase"
git fetch origin
git rebase origin/main
git stash pop

如果想让 Git 自动暂存未提交修改，可以用：

git pull --rebase --autostash origin main

注意：这条命令会影响当前 checkout 的分支。先确认自己在目标 feature 分支上。

push 时远程比当前要领先怎么办？

先区分清楚：准备开 PR 前，优先同步的是最新 main，不是盲目 git pull 当前分支。

推荐流程：

git status
git fetch origin
git rebase origin/main

这个习惯很重要。别人从其他分支合 PR 之后，main 可能已经领先；提前 rebase 到最新 main，通常能更早发现冲突，也能避免 PR 里混入过期历史。

如果已经在自己的分支上提交了 commit，再执行：

git rebase origin/main

Git 会把你的 commit 临时拿下来，再放到最新 main 后面重新应用。遇到冲突时按下面处理：

git status
# 手动解决冲突文件
git add <conflict-file>
git rebase --continue

如果发现 rebase 方向不对，直接取消：

git rebase --abort

当 rebase origin/main 之后，git status 仍然可能显示：

Your branch and 'origin/xxx' have diverged,
and have N and M different commits each, respectively.

这通常不是说你落后于 origin/main，而是说本地分支和远程同名分支 origin/xxx 分叉了。比如之前用同一个分支开过 PR，PR 已经合进 main，但远程功能分支还停在旧位置。

这时不要直接 git pull。git pull 会把 origin/xxx 拉回来整合，可能产生重复提交或无意义 merge。先确认当前分支已经基于最新 main：

git log --oneline --graph --decorate --left-right origin/main...HEAD

如果确认当前分支只是 origin/main + 你的新 commit，就更新远程功能分支：

git push --force-with-lease origin xxx

一句话记忆：

rebase origin/main：同步主线。
push --force-with-lease：更新远程功能分支。
不要用 git pull 去解决已经 rebase 到 main 后的同名分支分叉。

Flowra 工作流里 Git 怎么走？

这部分更像项目工具链笔记，主线可以按下面走：

flowra ws myWorkspace
git clone <repo-url>

进入 workspace 后修改代码，然后格式化：

flowra format your_project

再走正常 Git 流程：

git add -A
git commit -m "xxx"
git push

如果需要修改最后一次提交：

git commit --amend
git push --force-with-lease

pre-commit 应该什么时候跑？

安装：

pip install pre-commit

通常在 git add 后、git commit 前跑：

git add -A
pre-commit run --all-files
git commit -m "Describe your change"

如果 hook 自动修了文件，需要重新暂存：

git add -A
git commit -m "Describe your change"

怎么修改过去的 commit message？

如果历史 commit message 写错，并且 CI 或规范检查依赖 commit message，可以用 interactive rebase 的 reword。

选择要修改的最近 N 个 commit：

git rebase -i HEAD~N

在编辑器里，把需要修改 message 的那几行从 pick 改成 reword：

reword a1b2c3d feat: old message
pick   b2c3d4e fix: another message

保存退出后，Git 会逐个打开 commit message 编辑器。

如果这些 commit 已经推送到远端，修改完成后需要：

git push --force-with-lease

怎么合并过去的多个 commits？

先更新当前分支：

git checkout feat/op_matcher
git fetch origin
git pull --rebase

把最近 6 个提交放进 interactive rebase：

git rebase -i HEAD~6

如果历史里包含 merge commit，并且需要保留 merge 结构，可以改用：

git rebase -i --rebase-merges HEAD~6

编辑器里会看到类似：

pick a1b2c3d feat(op_matcher): reconstruct...
pick b2c3d4e feat(op_matcher): reconstruct...
pick c3d4e5f feat(op_matcher): reconstruct...
pick d4e5f6g feat(op_matcher): support bf16
pick e5f6g7h feat(op_matcher): reconstruct...
pick f6g7h8i feat(op_matcher): optimize compare

如果想把前 3 个合成 1 个，把后两行改成 squash 或 s：

pick a1b2c3d feat(op_matcher): reconstruct...
s    b2c3d4e feat(op_matcher): reconstruct...
s    c3d4e5f feat(op_matcher): reconstruct...
pick d4e5f6g feat(op_matcher): support bf16
pick e5f6g7h feat(op_matcher): reconstruct...
pick f6g7h8i feat(op_matcher): optimize compare

保存退出后，Git 会让你编辑新的 commit message。删掉重复内容，保留一个清晰的 summary 和必要 bullet 即可。

如果 rebase 过程中遇到冲突：

git add <conflict-file>
git rebase --continue

如果这些 commit 已经推送过，最后同步远端：

git push --force-with-lease

git commit -c HEAD 是什么？

git commit -c HEAD 会复用当前分支最后一次提交的 message 作为模板，并打开编辑器让你修改，然后创建一个新提交。

基本流程：

git add .
git commit -c HEAD

它和 git commit --amend 的区别：

  • git commit --amend：修改上一次提交，不产生额外的新提交，而是用新 commit 替换旧 commit。
  • git commit -c HEAD：创建一个全新的提交，只是借用上一次提交信息作为草稿。

-c 和 -C 的区别：

  • git commit -c HEAD：复用信息，并打开编辑器供你修改。
  • git commit -C HEAD：直接复用信息，不打开编辑器。

如果想“复制上个提交的消息并改一改，然后发一个新提交”，用 git commit -c HEAD。

合并冲突时怎么保留某一边？

如果你明确知道要保留当前分支这一边：

git checkout --ours op_matcher/verify_eau_opcode.py
git add op_matcher/verify_eau_opcode.py

如果你明确知道要保留另一边：

git checkout --theirs op_matcher/verify_eau_opcode.py
git add op_matcher/verify_eau_opcode.py

注意：在 rebase 场景里，ours 和 theirs 的直觉可能和 merge 场景不同。执行前最好用 git status 和冲突内容确认清楚。

怎么修正最近一次 commit 的 author？

如果最近一次提交的 author 信息不对：

git commit --amend --reset-author --no-edit

如果该提交已经推送过：

git push --force-with-lease

WSL 里 clone 项目失败可以先看什么？

如果是 DNS 或网络解析问题，可以先检查 /etc/resolv.conf：

sudo vim /etc/resolv.conf

有些网络环境里，需要把 nameserver 换成本机或可用 DNS 的 IP。这个属于环境问题，不一定是 Git 本身的问题。

amend 之后为什么普通 git push 会让我先 pull？

假设最开始本地和远端都是提交 A：

remote: A
local:  A

执行：

git commit --amend

本地提交会从 A 变成一个新的提交，比如 B：

remote: A
local:  B

B 不是追加在 A 后面的提交，而是替换了 A。这不是 fast-forward 关系，所以普通 git push 会拒绝，并提示你先 pull。

但这个场景通常不应该先 pull。你真正想做的是用本地新的 B 替换远端旧的 A，因此应该：

git push --force-with-lease

--force 和 --force-with-lease 有什么区别？

--force 会直接强推，不太管远端是否已经被别人更新。

--force-with-lease 会先检查远端是否还是你本地记录里的状态。如果远端在这期间被别人更新过，它会拒绝推送，避免你覆盖别人的提交。

所以平时更推荐：

git push --force-with-lease

怎么退回以前的提交？

如果提交已经推送到远端，最保守的撤销方式是 revert。它会新增一个“反向提交”，保留历史：

git revert HEAD
git push

如果后来决定把“原 commit + revert commit”都从历史中删掉，可以用 interactive rebase：

git rebase -i HEAD~3

在打开的编辑器里，把这两条从：

pick <sha1> feat: original change
pick <sha2> revert: revert original change

改成：

drop <sha1> feat: original change
drop <sha2> revert: revert original change

如果 rebase 过程中状态混乱，先退出：

git rebase --abort

rebase 成功后检查：

git status
git log --oneline -5

如果此时看到类似：

Your branch is behind origin/... by 2 commits

并且确认这 2 个 commit 正是刚刚 drop 掉的那两个，不要 git pull，否则会把它们又拉回来。应该执行：

git push --force-with-lease

完整链路：

git revert HEAD
git push

git rebase -i HEAD~3
# 把原 commit 和 revert commit 改成 drop

git status
git log --oneline -5

git push --force-with-lease

两个本地 clone 同一个分支，一个 force push 后另一个怎么同步？

场景：

  • 本地 1 对 feat/dump_llo 做了 rebase。
  • 本地 1 执行了 git push --force-with-lease。
  • 远端 origin/feat/dump_llo 变成一条新的提交历史。
  • 本地 2 仍然保留旧历史。

这时本地 2 执行：

git pull

可能会看到：

+ 815a613...2889e70 feat/dump_llo -> origin/feat/dump_llo  (forced update)
hint: You have divergent branches and need to specify how to reconcile them.
fatal: Need to specify how to reconcile divergent branches.

原因是本地旧历史和远端新历史已经不是 fast-forward 关系。git pull 不知道你想 merge、rebase，还是只接受 fast-forward。

如果本地 2 没有需要保留的提交，最直接的做法是让它完全对齐远端：

git fetch origin
git switch feat/dump_llo
git reset --hard origin/feat/dump_llo

如果担心本地 2 上有内容需要找回，先备份：

git branch backup/feat-dump-llo-local2
git fetch origin
git reset --hard origin/feat/dump_llo

这个场景不推荐直接 git pull，因为它会尝试整合两条不同历史，而你的目标通常只是“本地 2 跟远端同步”。

Git worktree 和 branch 有什么区别？

branch 是 Git 历史上的提交指针，例如 main、codex/reconstruct。它记录代码历史走到哪里。

worktree 是磁盘上的实际工作目录。一个仓库可以有多个 worktree，每个 worktree 都是一份可以编辑、运行、提交代码的工作现场。

可以这样理解：

repo
├─ main worktree -> main
├─ worktree A    -> feature/login
└─ worktree B    -> codex/reconstruct

关键点：

  • 一个 worktree 同一时间通常只能 checkout 一个 branch，或者处于 detached HEAD。
  • 一个 branch 同一时间通常只能被一个 worktree 占用。
  • 如果某个 branch 已经被别的 worktree 使用，新的 worktree 不能直接 git switch 到这个 branch。
  • 删除 Codex 对话不等于删除 Git worktree。worktree 是 Git 记录的本地目录，需要用 Git 命令清理。

worktree 常用命令有哪些？

查看当前有哪些 worktree：

git worktree list

查看某个 worktree 的状态：

git -C "C:/path/to/worktree" status --short --branch

新建一个 worktree，并让它绑定新分支：

git worktree add "C:/path/to/new-worktree" -b feat/example main

新建一个 worktree，并让它 checkout 已有分支：

git worktree add "C:/path/to/new-worktree" feat/example

删除不再使用的普通 worktree：

git worktree remove "C:/path/to/old-worktree"

如果目录已经被手动删掉，但 Git 记录还在：

git worktree prune

查看当前 worktree 绑定的分支：

git branch --show-current

如果显示不出分支名，可能是 detached HEAD：

git status --short --branch

切换分支：

git switch main
git switch codex/reconstruct

worktree 里目标分支被占用怎么办？

先查看 worktree 结构：

git worktree list

可能会看到类似：

E:/my_github_page/wqh011128.github.io                         [codex/reconstruct]
C:/Users/吴启航/.codex/worktrees/4855/wqh011128.github.io     (detached HEAD)

这说明 codex/reconstruct 分支已经被 E:/my_github_page/wqh011128.github.io 这个 worktree 占用，当前 Codex worktree 不能直接接管这个分支。

如果被占用的是主工作目录，不能用下面这条命令删除：

git worktree remove "E:/my_github_page/wqh011128.github.io"

否则会看到：

fatal: 'E:/my_github_page/wqh011128.github.io' is a main working tree

正确思路是让占用分支的 worktree 先切到别的分支，例如 main，从而释放目标分支。

推荐步骤：

git worktree list
git -C "E:/my_github_page/wqh011128.github.io" status --short --branch
git -C "C:/Users/吴启航/.codex/worktrees/4855/wqh011128.github.io" status --short --branch

确认没有未提交内容后，让主工作目录切回 main：

git -C "E:/my_github_page/wqh011128.github.io" switch main

然后在当前 worktree 接管目标分支：

git switch codex/reconstruct

最后确认：

git worktree list
git status --short --branch

期望结构类似：

E:/my_github_page/wqh011128.github.io                         [main]
C:/Users/吴启航/.codex/worktrees/4855/wqh011128.github.io     [codex/reconstruct]

worktree 上有本地修改时怎么处理？

如果目标分支或目标 worktree 上有本地修改，先决定是否保留。不要直接切分支、删除 worktree 或 reset。

如果确认不保留已跟踪文件的修改：

git restore --staged --worktree .

如果还存在未跟踪文件或目录，例如 skills/：

git clean -fd -- "skills/"

然后再按需要 rebase 到 main：

git rebase main

日常使用 worktree 的习惯是什么？

  • 主目录长期放 main。
  • 每个 Codex 任务使用单独 worktree 和单独分支。
  • 一个长期开发分支尽量固定在一个 worktree 使用。
  • 用完的 worktree 先检查 git status --short --branch，确认干净后再 git worktree remove。
  • 看到 detached HEAD 先不要慌，它只是说明当前 worktree 没有绑定分支；先用 git worktree list 判断目标分支是否被别的 worktree 占用。
]]>