models/README.md

# Parallel DMTD：训练与测试逻辑

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

```text
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` 划分：

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

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

cycle head 由以下规则确定：

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

每个 cycle 的 shadow 输入为：

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

因此完整 shadow 序列为：

```text
[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：

```text
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 序列：

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

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

对于位置 `i`：

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

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

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

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

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

训练过程：

```text
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

真实序列：

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

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

```text
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 序列：

```text
[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 顺序。

可见性为：

```python
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
```

历史模式为：

```text
[real,real,real,real] ... [head,MASK,MASK,MASK]
```

即：

```text
[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 规则：

```python
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：

```text
[x_c,M_1,M_2,M_3]
```

其内部 attention 为：

```text
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 的可见性为：

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

因此：

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

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

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

```text
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：

```text
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：

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

每处理一个真实 token：

```text
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 只运行一次：

```text
[x_c,M,M,M] → P28
```

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

```text
[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：

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

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

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

```text
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

将：

```text
[x4,x5,x6,x7]
```

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

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

```text
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)
```

完成后：

```text
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：

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

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

```text
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]`；随后只有：

```text
embed(x8) + s8
```

进入 S8，产生 `x9`。

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

#### Refresh real queries：前四个位置

```text
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 都可以看到：

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

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

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

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

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

```text
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：

```python
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 语义：

```text
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

构造：

```text
[x8,MASK9,MASK10,MASK11]
```

依次通过 Layer 0 到 Layer 27。

对于每一个 parallel layer `l`：

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

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

当前层 attention 使用：

```text
[persistent real history ; current temporary shadow block]
```

当前层计算结束后：

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

完成 Layer 27 后，只保存：

```text
[s8,s9,s10,s11]
```

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

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

### 7.5 Causal Refresh 与 Bidirectional Refresh

`Causal-Parallel-Refresh` 当前 local shadow block：

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

`Bidirectional-Parallel-Refresh` 当前 local shadow block：

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

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

### 7.6 Cycle 内 S8 decoding

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

S8 的每个 Layer 28~35 分别追加真实位置 `8~11` 的 K/V。

Cycle 结束后：

```text
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：

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

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

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

```text
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

构造：

```text
[x8,MASK9,MASK10,MASK11]
```

对于每一个 parallel layer `l`：

```text
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：

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

因此 forward 后：

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

同时保存 Layer 27 最终输出：

```text
[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 完全相同：

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

Cycle 结束后：

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

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

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

下一个 cycle 直接计算：

```text
[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

```text
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

```text
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
```

最关键的区别：

```text
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

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

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

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

### Step B：Bidirectional shadow forward

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

每层：

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

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

最终只保存：

```text
[s8,s9,s10,s11]
```

### Step C：Sequential decoding

```text
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

释放：

```text
[s8,s9,s10,s11]
```

然后：

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

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