xuan-luo/parallel-eval/models - Storage Bucket

Parallel DMTD：训练与测试逻辑

四个模型使用相同的基础结构：

Qwen3-4B
├── 36 个 Transformer layers
├── Parallel layers:   P28 = layers[0:28]
├── Sequential layers: S8  = layers[28:36]
├── Cycle length:      τ   = 4
└── MASK token id:         = 151660

四个模型由两个正交选项组成：

模型	当前 cycle 内 P28 shadow attention	历史 cycle 表示
Causal-Parallel-Norefresh	Causal	保留历史 shadow KV
Causal-Parallel-Refresh	Causal	使用历史 real KV，不保留 shadow KV
Bidirectional-Parallel-Norefresh	Bidirectional	保留历史 shadow KV
Bidirectional-Parallel-Refresh	Bidirectional	使用历史 real KV，不保留 shadow KV

后文使用：

• x_c：当前 cycle 的第一个真实 token，即 cycle head；
• M：MASK token；
• s_0,s_1,s_2,s_3：当前 cycle 经过 P28 后得到的四个最终 shadow hidden states；
• C_P[l]：第 l 个 parallel layer 的 KV cache，l=0,...,27；
• C_S[l]：第 l 个 sequential layer 的 KV cache，l=28,...,35。

---

一、训练逻辑

1. 共同训练逻辑

1.1 Cycle 划分

训练时将序列按照 τ=4 划分：

真实序列：
[x0, x1, x2, x3] [x4, x5, x6, x7] ...

cycle id：
[0,  0,  0,  0 ] [1,  1,  1,  1 ] ...

cycle head 由以下规则确定：

is_cycle_head = position_ids % mtp_horizon == 0
cycle_id = position_ids // mtp_horizon

每个 cycle 的 shadow 输入为：

[真实 cycle head, MASK, MASK, MASK]

因此完整 shadow 序列为：

[x0, M, M, M] [x4, M, M, M] ...

1.2 P28 输出与真实 embedding 合并

P28 处理 shadow 序列，产生每个位置的 parallel hidden：

[ h_i^{parallel}=P28(shadow)_i ]

S8 的输入为：

[ h_i^{S8}=Embed(x_i)+h_i^{parallel} ]

两个分量的含义不同：

• Embed(x_i) 是当前位置的真实 token；
• h_i^{parallel} 是 cycle 开始时可以提前计算的 shadow/lookahead 表示。

例如一个 cycle：

P28([x_c,M,M,M]) → [s_0,s_1,s_2,s_3]

S8 的训练输入：
embed(x_c)     + s_0
embed(x_{c+1}) + s_1
embed(x_{c+2}) + s_2
embed(x_{c+3}) + s_3

训练时真实序列全部已知，因此可以在同一个 teacher-forced forward 中计算四个位置的 S8 输入。

1.3 Loss

四个模型都使用普通 next-token causal language modeling loss：

[ \mathcal{L} =-\sum_i \log p(x_{i+1}\mid x_{\le i},h_i^{parallel}) ]

没有额外 LM head，也没有额外 multi-token loss。区别只在 P28 shadow representation 如何构造。

1.4 Packed sequence

使用 packing 时：

• 每个 document 的 position_ids 从 0 重新开始；
• cycle 不能跨越 document boundary；
• real 和 shadow attention 都不能跨 document；
• 同一 batch 应具有兼容的 cycle layout。

---

2. Causal-Parallel-Norefresh 训练

Shadow 序列：

[x0,M,M,M] [x4,M,M,M] [x8,M,M,M] ...

整个 shadow 序列通过 P28，使用标准 causal attention。

对于位置 i：

shadow position i 只能看到 shadow positions <= i

因此当前 cycle 可以看到历史 cycle 的 shadow 表示：

历史：
[x0,M,M,M] [x4,M,M,M]

当前：
[x8,M,M,M]

历史非 head 位置始终是 MASK representation，不会被真实 token 替换。

训练过程：

shadow_ids
    │
    └── P28 causal shadow forward
              │
              └── parallel_hidden

inputs_embeds + parallel_hidden
              │
              └── S8 causal forward
                        │
                        └── LM head + next-token loss

Norefresh 的关键训练语义是：

后续 cycle 会读取之前 cycle 的 shadow K/V，因此历史 shadow representation 是模型 上下文的一部分。

---

3. Causal-Parallel-Refresh 训练

训练时同时运行共享权重的 real branch 和 shadow branch。

3.1 Real branch

真实序列：

[x0,x1,x2,x3] [x4,x5,x6,x7] ...

通过 P28 标准 causal forward。每个 parallel layer 都产生自己的 real K/V：

Layer 0 : K_real[0], V_real[0]
Layer 1 : K_real[1], V_real[1]
...
Layer 27: K_real[27], V_real[27]

Real branch 的最终 hidden states 不直接进入 S8。它只为 shadow branch 提供历史 cycle 的真实 K/V。

3.2 Shadow branch

Shadow 序列：

[x0,M,M,M] [x4,M,M,M] ...

在每个 P28 layer，shadow query 读取：

1. 严格更早 cycle 的 real K/V；
2. 当前 cycle 的 shadow K/V；
3. 当前 cycle 内遵守 causal 顺序。

可见性为：

real_visible = real_kv and real_cycle < query_cycle
shadow_visible = shadow_kv and shadow_cycle == query_cycle and shadow_position <= query_position
visible = real_visible or shadow_visible

历史模式为：

[real,real,real,real] ... [head,MASK,MASK,MASK]

即：

[1111] [1111] [1000]

这在一次 teacher-forced forward 中模拟测试时：

1. 上一 cycle 已完成；
2. 上一 cycle 用真实 token causal P28 forward 完成 refresh；
3. 当前 cycle 再执行 shadow forward。

Refresh 训练的关键语义是：

Shadow branch 只使用当前 cycle 的临时 shadow K/V；跨 cycle 历史使用 real K/V， 不使用历史 shadow K/V。

---

4. Bidirectional-Parallel-Norefresh 训练

它与 Causal-Parallel-Norefresh 的区别仅在当前 cycle 内的 shadow attention。

Attention 规则：

same_cycle = query_cycle == key_cycle
causal_history = key_position <= query_position
visible = same_cycle or causal_history

含义：

• 跨 cycle 保持 causal；
• 当前 cycle 的四个 shadow 槽位完全双向可见；
• 历史 cycle 仍然是 [head,MASK,MASK,MASK] shadow representation。

当前 cycle：

[x_c,M_1,M_2,M_3]

其内部 attention 为：

x_c ↔ M_1 ↔ M_2 ↔ M_3

这里双向可见的只是 MASK/shadow representation，不包含未来真实 token。

Norefresh 训练仍然要求历史 shadow representation 被后续 cycle 使用。

---

5. Bidirectional-Parallel-Refresh 训练

它结合：

1. 历史 cycle 使用 real K/V；
2. 当前 cycle 的 shadow block 使用 bidirectional attention。

每个 P28 layer 的可见性为：

real_visible = real_kv and real_cycle < query_cycle
shadow_visible = shadow_kv and shadow_cycle == query_cycle
visible = real_visible or shadow_visible

因此：

历史 cycle：
[real,real,real,real] [real,real,real,real]

当前 cycle：
[head,MASK,MASK,MASK]，四个位置双向可见

Real branch 仍然只提供每层历史 K/V；最终送入 S8 的是 shadow branch 最后一层输出：

parallel_hidden = shadow_hidden_after_layer_27
h_S8 = inputs_embeds + parallel_hidden

Refresh 训练不要求历史 shadow KV 跨 cycle 保留。

---

二、测试逻辑

6. 共同测试结构

测试仍然使用 DMTD-style cycled decoding。

每个 cycle 的第一个真实 token 已知后：

1. 根据模型类型处理历史 P28 KV；
2. 一次性计算当前 cycle 的四个 shadow hidden；
3. 保存最终 [s_0,s_1,s_2,s_3]；
4. cycle 内逐 token 运行 S8；
5. S8 每一步使用一个真实 token embedding 和对应 shadow hidden。

当前 cycle：

P28([x_c,M,M,M]) → [s_0,s_1,s_2,s_3]

embed(x_c)     + s_0 → S8 → x_{c+1}
embed(x_{c+1}) + s_1 → S8 → x_{c+2}
embed(x_{c+2}) + s_2 → S8 → x_{c+3}
embed(x_{c+3}) + s_3 → S8 → x_{c+4}

x_{c+4} 是下一个 cycle head。

6.1 S8 cache：四个模型完全相同

后 8 层中的每一层都保存自己的 real KV：

C_S[28], C_S[29], ..., C_S[35]

每处理一个真实 token：

Layer 28 append 当前 token 的 K/V
Layer 29 append 当前 token 的 K/V
...
Layer 35 append 当前 token 的 K/V

S8 cache：

• 每个真实 token 增加一个位置；
• 始终是 dense real-token history；
• 永远不保存 MASK；
• 永远不 refresh；
• 永远不 crop。

6.2 当前 cycle shadow hidden

每个 cycle 只运行一次：

[x_c,M,M,M] → P28

只需要为 cycle 内 S8 decoding 持久保存：

[s_0,s_1,s_2,s_3]

这四个 hidden 是 Layer 27 输出后的最终 parallel hidden。

P28 每一层在计算当前 block attention 时仍然需要当前 block 的临时 Q/K/V，但是否将 这些 K/V 写入跨 cycle cache，取决于 Refresh 或 Norefresh。

---

7. Refresh 测试逻辑

适用于：

• Causal-Parallel-Refresh
• Bidirectional-Parallel-Refresh

7.1 Persistent P28 cache

Refresh 模型持久保存 28 组 real KV cache：

C_P[0], C_P[1], ..., C_P[27]

其中每一组都只包含已经 refresh 的真实 token K/V。

例如进入位置 8~11 的 cycle 前，P28 real cache 只包含位置 0~3：

Layer 0 : [R0 R1 R2 R3]
Layer 1 : [R0 R1 R2 R3]
...
Layer 27: [R0 R1 R2 R3]

上一 cycle 的真实 token x4,x5,x6,x7 已经生成，但还未进入 P28 persistent cache。

7.2 Refresh 上一 cycle

将：

[x4,x5,x6,x7]

作为一个真实 block，通过全部 28 个 parallel layers，使用标准 causal attention。

在每一层分别追加该层产生的 real K/V：

Layer 0 append K_real[0](x4:x7), V_real[0](x4:x7)
Layer 1 append K_real[1](x4:x7), V_real[1](x4:x7)
...
Layer 27 append K_real[27](x4:x7), V_real[27](x4:x7)

完成后：

Layer 0 : [R0 ... R7]
Layer 1 : [R0 ... R7]
...
Layer 27: [R0 ... R7]

因为 Refresh 模式从未将上一 cycle 的 shadow K/V 写入 persistent cache，所以：

• 不需要 crop shadow KV；
• 不需要替换 shadow KV；
• 直接 append real KV 即可。

7.3 将 refresh 与当前 shadow 合并为一次 P28 forward

上面的“refresh 上一 cycle”和“forward 当前 shadow block”描述的是两个逻辑阶段。 实际编写 inference 时，应将它们合并成一次带定制 attention mask 的 P28 forward， 避免先后启动两次 4-token forward。

在稳态 cycle 中，P28 一次接收 8 个 token slots：

P28 input:
[x4,x5,x6,x7 | x8,MASK9,MASK10,MASK11]
 └ refresh real ┘  └── current shadow ──┘
       4 tokens              4 tokens

所以 cycle 第一次 forward 的计算粒度是：

Layers 0~27 (P28): 8 个 token slots
Layers 28~35 (S8): 仅当前 cycle head x8，共 1 个 token

不能把 8 个 token 全部送入 S8。前四个 token 仅用于生成 P28 real KV，当前 shadow 四槽仅用于生成 [s8,s9,s10,s11]；随后只有：

embed(x8) + s8

进入 S8，产生 x9。

对于每一个 parallel layer l=0,...,27，一次性计算这 8 个位置的 Q/K/V，并使用以下 block attention mask。

Refresh real queries：前四个位置

query x4 只能看到 persistent real history 和 x4
query x5 只能看到 persistent real history 和 x4:x5
query x6 只能看到 persistent real history 和 x4:x6
query x7 只能看到 persistent real history 和 x4:x7

它们使用标准 causal attention，并且绝不能看到当前 cycle 的 shadow 位置 8~11。

Current shadow queries：后四个位置

所有当前 shadow query 都可以看到：

1. P28 persistent real history；
2. 本次 forward 中刚计算的 refresh real positions x4:x7；
3. 当前 cycle 的 local shadow block。

当前 shadow block 内部的可见性由模型决定：

Causal-Parallel-Refresh:
    当前四槽内部使用 causal mask

Bidirectional-Parallel-Refresh:
    当前四槽内部完全双向可见

以 Bidirectional-Parallel-Refresh 为例，单层 mask 的逻辑结构是：

Keys →
Queries ↓        old real cache    refresh real 4    current shadow 4

refresh real 4       causal            causal               ×
current shadow 4       ✓                 ✓            bidirectional

其中 × 表示 refresh real queries 禁止读取当前 shadow。

每层 forward 后必须拆分 KV：

real_kv = kv[:, :, :4, :]       # x4:x7
shadow_kv = kv[:, :, 4:8, :]    # x8:MASK11

append_to_persistent_cache(real_kv)
discard_after_this_layer(shadow_kv)

也就是说：

• 前四个 refresh real K/V 写入该层 C_P[l]；
• 后四个 current shadow K/V 只参与本层 attention，不写入 persistent cache；
• real hidden 和 shadow hidden 都传给下一个 P28 layer；
• Layer 27 完成后，只持久保存最后四个 shadow outputs： [s8,s9,s10,s11]。

因此，合并 forward 不改变 Refresh 的 cache 语义：

P28 persistent cache 只增加 4 个 real KV positions，
不会增加当前 cycle 的 4 个 shadow KV positions。

第一个生成 cycle 是特殊情况：没有上一生成 cycle 可以 refresh，因此第一次 P28 forward 只有当前 [head,MASK,MASK,MASK] 四个 shadow slots。之后的稳态 cycle 才使用 上述 8-token P28 forward。

7.4 Forward 当前 shadow block

构造：

[x8,MASK9,MASK10,MASK11]

依次通过 Layer 0 到 Layer 27。

对于每一个 parallel layer l：

Persistent history:
C_P[l] = real K/V positions 0~7

Current local block:
K_shadow[l], V_shadow[l] positions 8~11

当前层 attention 使用：

[persistent real history ; current temporary shadow block]

当前层计算结束后：

• shadow K/V 不写入 C_P[l]；
• shadow K/V 只在当前层 forward 中临时存在；
• 当前层输出 hidden 传给下一层。

完成 Layer 27 后，只保存：

[s8,s9,s10,s11]

除本次合并 forward 中已经追加的四个 refresh real positions 外，当前 shadow block 不会继续改变 P28 persistent cache：

Layer 0 : [R0 ... R7]
...
Layer 27: [R0 ... R7]

7.5 Causal Refresh 与 Bidirectional Refresh

Causal-Parallel-Refresh 当前 local shadow block：

s8 看到 shadow positions <= 8
s9 看到 shadow positions <= 9
s10 看到 shadow positions <= 10
s11 看到 shadow positions <= 11

Bidirectional-Parallel-Refresh 当前 local shadow block：

s8,s9,s10,s11 相互全部可见

两者都读取相同的 persistent real P28 history，且都不保存当前 shadow KV。

7.6 Cycle 内 S8 decoding

embed(x8)  + s8  → S8 → x9
embed(x9)  + s9  → S8 → x10
embed(x10) + s10 → S8 → x11
embed(x11) + s11 → S8 → x12

S8 的每个 Layer 2835 分别追加真实位置 `811` 的 K/V。

Cycle 结束后：

P28 persistent cache:
    每层只包含 real positions 0~7

S8 persistent cache:
    每层包含 real positions 0~11

Saved shadow hidden:
    cycle 结束后可以释放

下一个 cycle 以 x12 开始时：

1. 将真实 [x8,x9,x10,x11] causal forward P28；
2. P28 每层 append real KV positions 8~11；
3. 再计算 [x12,MASK13,MASK14,MASK15] 的当前 shadow hidden；
4. 当前 shadow K/V 仍然不持久保存。

Refresh 模式的核心不变量：

P28 persistent cache 永远只包含 real KV；当前 cycle shadow KV 永远是临时值。

---

8. Norefresh 测试逻辑

适用于：

• Causal-Parallel-Norefresh
• Bidirectional-Parallel-Norefresh

8.1 Persistent P28 cache

Norefresh 不会在下一 cycle 使用真实 token 重算 P28，因此历史 context 只能来自之前 cycle 的 shadow representation。

所以必须持久保存 28 组 shadow KV cache：

C_P[0], C_P[1], ..., C_P[27]

每个 parallel layer 都保存自己的 shadow K/V，不能只保存 Layer 27。

例如进入位置 8~11 的 cycle 前：

Layer 0 : [S0 S1 S2 S3 | S4 S5 S6 S7]
Layer 1 : [S0 S1 S2 S3 | S4 S5 S6 S7]
...
Layer 27: [S0 S1 S2 S3 | S4 S5 S6 S7]

这里每个 S 都是相应 layer 在之前 shadow forward 中产生的 K/V。

8.2 Forward 当前 shadow block

构造：

[x8,MASK9,MASK10,MASK11]

对于每一个 parallel layer l：

Persistent history:
C_P[l] = previous cycles' shadow K/V

Current local block:
K_shadow[l], V_shadow[l] positions 8~11

当前层 attention 使用历史 shadow cache 和当前 local block。计算完成后，将当前 block 的 四个 shadow K/V 追加到该层 persistent cache：

Layer 0 append S8:S11
Layer 1 append S8:S11
...
Layer 27 append S8:S11

因此 forward 后：

Layer 0 : [S0 ... S11]
Layer 1 : [S0 ... S11]
...
Layer 27: [S0 ... S11]

同时保存 Layer 27 最终输出：

[s8,s9,s10,s11]

区别是：

• 28 层 shadow K/V 用于后续 cycle 的 P28 attention；
• 最终四个 shadow hidden 用于当前 cycle 的 S8 decoding。

两者用途不同，Norefresh 中都需要保存。

8.3 Causal Norefresh 与 Bidirectional Norefresh

Causal-Parallel-Norefresh：

• 历史 shadow cache 全部可因果读取；
• 当前 cycle 内也是 causal。

Bidirectional-Parallel-Norefresh：

• 历史 cycle 作为已有 shadow cache 被当前 cycle 读取；
• 当前 cycle 的四个 local shadow 槽位双向可见；
• 当前 cycle 计算出的 bidirectional shadow K/V 会被持久保存；
• 下一 cycle 会读取这些历史 shadow K/V。

8.4 Cycle 内 S8 decoding

与 Refresh 完全相同：

embed(x8)  + s8  → S8 → x9
embed(x9)  + s9  → S8 → x10
embed(x10) + s10 → S8 → x11
embed(x11) + s11 → S8 → x12

Cycle 结束后：

P28 persistent cache:
    每层包含 shadow positions 0~11

S8 persistent cache:
    每层包含 real positions 0~11

Saved shadow hidden:
    cycle 结束后可以释放

下一个 cycle 直接计算：

[x12,MASK13,MASK14,MASK15] → P28

它会读取已经保存的历史 shadow KV，不执行 real causal refill。

Norefresh 模式的核心不变量：

P28 persistent cache 保存所有历史 shadow KV；这些 KV 是后续 cycle 的 P28 context。

---

9. Refresh 与 Norefresh 的 cache 对比

Refresh

P28 cache:
    28 层，每层只保存 real KV
    每个 cycle 开始时，将上一 cycle 的真实 token causal forward P28 后 append
    当前 cycle shadow KV 不写入 cache

Current cycle:
    只保存 P28 最终输出的 4 个 shadow hidden

S8 cache:
    8 层，每层保存所有真实 token KV

Norefresh

P28 cache:
    28 层，每层保存历史 shadow KV
    每个 cycle 的 shadow forward 都向每层 append 4 个 shadow KV
    不执行 real causal refill

Current cycle:
    保存 P28 最终输出的 4 个 shadow hidden

S8 cache:
    8 层，每层保存所有真实 token KV

最关键的区别：

Refresh:
    shadow KV 是当前 P28 forward 的临时张量
    跨 cycle 保存的是 real KV

Norefresh:
    shadow KV 是后续 cycle 的历史 context
    必须跨 cycle 保存

---

10. 一个完整的 Bidirectional-Parallel-Refresh cycle

假设：

• 上一 cycle：x4,x5,x6,x7
• 当前 cycle：位置 8~11
• 当前 head：x8

Step A：Refresh

[x4,x5,x6,x7] → causal P28

每个 Layer 0~27 append 自己的 real K/V：

C_P[0]  append R4:R7
C_P[1]  append R4:R7
...
C_P[27] append R4:R7

Step B：Bidirectional shadow forward

[x8,M9,M10,M11] → bidirectional-within-cycle P28

每层：

attention history = 该层 persistent real K/V
attention current = 该层 temporary shadow K/V

当前四个 shadow 位置双向可见，但 temporary shadow K/V 不写入 persistent cache。

最终只保存：

[s8,s9,s10,s11]

Step C：Sequential decoding

embed(x8)  + s8  → S8 → x9
embed(x9)  + s9  → S8 → x10
embed(x10) + s10 → S8 → x11
embed(x11) + s11 → S8 → x12

每次 S8 forward 都向 Layer 28~35 的 cache 各追加一个真实 K/V。

Step D：进入下一 cycle

释放：

[s8,s9,s10,s11]

然后：

[x8,x9,x10,x11] → causal P28 refresh
[x12,M13,M14,M15] → bidirectional shadow P28

整个过程中，S8:S11 的 bidirectional shadow K/V 从未写入 persistent P28 cache。