ds4.c 学习笔记

Appearance

Part 4: 练习 ​

1. RoPE 与推理循环 ​

练习 1：RoPE 旋转计算 ​

题目：位置 pos=2, head_dim=4, n_rot=2, freq_base=10000。计算第 0 对维度的旋转角度。

提示：theta = pos × freq_base^(-2×0/n_rot) = pos × freq_base^0 = pos

参考答案（先自己尝试，再展开查看）
第 0 对维度（i=0）:
  theta = 2 × 10000^(−0/2) = 2 × 1 = 2

  cos(2) ≈ -0.416
  sin(2) ≈  0.909

  如果 tail[0] = 1.0, tail[1] = 0.0:
    tail[0]' = 1.0 × (-0.416) - 0.0 × 0.909 = -0.416
    tail[1]' = 1.0 × 0.909 + 0.0 × (-0.416) = 0.909

  向量从 (1, 0) 旋转到 (-0.416, 0.909) ≈ 旋转了 2 弧度 (114.6°)

练习 2：Prefill vs Decode 性能 ​

题目：MacBook Pro M3 Max 上，ds4 的 prefill 速度是 250 t/s，decode 速度是 27 t/s。为什么 decode 慢这么多？

参考答案（先自己尝试，再展开查看）
  1. Prefill 可以批量：所有 prompt token 同时通过每一层，GPU 可以并行计算多个 token 的矩阵乘法
  2. Decode 是串行的：每次只处理一个 token，GPU 利用率低
  3. 内存带宽瓶颈：decode 时每步需要读取大量权重（81GB 模型），但只计算一个 token。计算/带宽比极低
  4. KV cache 访问：decode 需要读取所有 KV cache 行做注意力，随着生成进行，KV cache 越来越大

简单说：prefill 是"计算密集"（GPU 擅长），decode 是"带宽密集"（GPU 的瓶颈）。

练习 3：自回归生成模拟 ​

题目：假设 logits 经过 softmax 后的概率分布为：

Token  概率
"cat"  0.5
"dog"  0.3
"the"  0.15
"is"   0.05

用贪心解码（argmax），模拟生成 "the cat is" 的过程（假设前两步已经确定了 "the cat"）。

参考答案（先自己尝试，再展开查看）
Step 0: 输入 "the"
  logits → softmax → argmax → "cat" (0.5)
  输出: "cat"

Step 1: 输入 "cat"（KV cache 已有 "the"）
  logits → softmax → argmax → "is" (0.5)
  输出: "is"

完整生成: "the" → "cat" → "is"

关键点：每一步的输入是上一步的输出，这就是"自回归"的含义。如果某一步的概率分布不同，后续所有步骤都会改变。

练习 4：推测解码效率 ​

题目：假设 MTP 模型的草稿接受率为 70%。每次 MTP 推测 2 个 token。

  1. 平均每次验证接受多少 token？
  2. 相比不用推测解码，加速多少？（假设验证一次的成本 = 正常 decode 一步）
参考答案（先自己尝试，再展开查看）
每次推测 2 个 token 的验证:

情况1: 前 2 个都正确 (70% × 70% = 49%)
  → 接受 2 个 token + 验证产生 1 个额外 token = 3 个 token
  → 成本: 1 次验证

情况2: 第 1 个正确，第 2 个错误 (70% × 30% = 21%)
  → 接受 1 个 token + 验证产生 1 个额外 token = 2 个 token
  → 成本: 1 次验证

情况3: 第 1 个就错误 (30%)
  → 接受 0 个 token + 验证产生 1 个 token = 1 个 token
  → 成本: 1 次验证

平均每次验证产出:
  0.49 × 3 + 0.21 × 2 + 0.30 × 1 = 1.47 + 0.42 + 0.30 = 2.19 token

不用推测: 每步 1 token
加速比: 2.19 / 1 = 2.19×（理论值）

实际加速更低：MTP 推测本身也有成本，且需要额外的 MTP 模型权重。

今日学习检查清单 ​

  • 能解释 RoPE 通过旋转编码位置的原理
  • 理解 YaRN 扩展如何支持超长上下文
  • 能区分 prefill 和 decode 阶段
  • 理解自回归生成的串行本质
  • 能描述完整的推理流程（从输入到输出）
  • 理解推测解码（MTP）的加速原理
  • 理解 decode 为什么比 prefill 慢

延伸挑战 ​

挑战 1（中级）：YaRN 扩展的可视化 ​

画一张图：x 轴是维度对编号（0 到 64），y 轴是位置（0 到 1M）。用不同颜色标注"外推区域"和"内插区域"。标注训练长度 65536 和 scale factor 16 的关系。

挑战 2（高级）：MTP 接受率分析 ​

阅读 ds4.c 中 MTP 推测解码的实现（行 16071-16320）。解释 mtp_margin 阈值的作用：在什么条件下 MTP 草稿会被丢弃？如果 MTP 草稿的准确率只有 50%，MTP 是加速还是减速？用数学推导。

2. 线程池 ​

练习 1：线程池行分配 ​

题目：8 个 token，3 个线程。计算每个线程的行范围。

参考答案（先自己尝试，再展开查看）
rows_per_thread = (8 + 3 - 1) / 3 = 10 / 3 = 3

Thread 0 (主线程): [0, 3)   → token 0, 1, 2
Thread 1:          [3, 6)   → token 3, 4, 5
Thread 2:          [6, 8)   → token 6, 7  (只有 2 个)

练习 2：为什么 wait 用 while 不用 if ​

题目：解释为什么 pthread_cond_wait 必须在 while 循环中。

参考答案（先自己尝试，再展开查看）
  1. 虚假唤醒：POSIX 允许 cond_wait 在没有 signal/broadcast 的情况下返回
  2. 竞态条件：多个 worker 被唤醒但只有一个能获取任务
  3. generation 检查：旧的 broadcast 可能在新任务设置前唤醒 worker
错误:
if (done < n_workers) cond_wait(...)   // 可能虚假唤醒后继续

正确:
while (done < n_workers) cond_wait(...)  // 每次醒来都重新检查

练习 3：并行退化的场景 ​

题目：g_parallel_depth 防止嵌套并行。举一个具体场景说明为什么需要这个保护。

参考答案（先自己尝试，再展开查看）
1. 主线程调用 ds4_parallel_for(1000, layer_ffn_batch, ...)
2. 所有线程开始执行 layer_ffn_batch
3. layer_ffn_batch 内部调用 ds4_parallel_for(256, swiglu_worker, ...)
4. 如果不阻止：
   - 每个线程又创建 N 个任务 → N² 个任务
   - 但只有 N 个线程 → 大量等待和同步开销
5. g_parallel_depth > 0 时退化为串行：
   - swiglu_worker 在当前线程内顺序执行
   - 避免嵌套线程同步的开销

练习 4：线程交错执行追踪 ​

题目：2 个 worker 线程 + 主线程，generation 从 10 开始。主线程调用 ds4_parallel_for(7, fn, ctx) 执行 n_rows=7 的任务。

要求：

  1. 计算每个线程的行范围
  2. 给定 Worker 0 先完成、Worker 1 后完成的顺序，追踪 done 值和 done_cond signal 的时机
  3. 画出完整的时序图
参考答案（先自己尝试，再展开查看）
rows_per_thread = (7 + 3 - 1) / 3 = 3
Thread 0 (主):  [0, 3)
Thread 1:       [3, 6)
Thread 2:       [6, 7)

完成顺序追踪:
1. Worker 0 完成 → __sync_fetch_and_add(&done, 1) → done=1, 1≠2, 不 signal
2. Worker 1 完成 → __sync_fetch_and_add(&done, 1) → done=2, 2==n_workers, signal done_cond
3. 主线程在 cond_wait 中被唤醒 → 检查 done >= n_workers → 条件满足 → 继续

时序图:
主线程              Worker 0           Worker 1
  │                    │                   │
  ├─ generation=11 ────┤                   │
  ├─ broadcast ───────→│ ← cond_wait      │ ← cond_wait
  │                    │                   │
  ├─ fn(ctx,0,3)      ├─ fn(ctx,3,6)      ├─ fn(ctx,6,7)
  │                    │                   │
  │                    ├─ done=1           │
  │                    │   (1≠2, 不signal) │
  │                    │                   ├─ done=2
  │ ←─────────────────── signal done_cond ─┤
  │                    │                   │
  ▼                    ▼                   ▼

练习 5：generation 竞态分析 ​

题目：假设去掉 generation 机制，改为只用一个 bool has_work 标志。给出一个具体的交错序列，展示在没有 generation 的情况下 worker 会如何"偷懒"（拿到旧任务或丢失新任务）。

参考答案（先自己尝试，再展开查看）
场景: 两次快速连续的 parallel_for 调用

步骤:
1. 主线程完成第一次 parallel_for，done=2
2. Worker 0 还没回到 cond_wait（还在循环顶部）
3. 主线程发起第二次 parallel_for:
   - has_work = true
   - done = 0
   - broadcast
4. Worker 0 恰好执行到检查 has_work → true
5. Worker 0 拿到第二次的任务参数
6. Worker 1 从 cond_wait 醒来，但 has_work 仍为 true
7. Worker 1 也拿到第二次的任务 → 重复执行！

对比: 有 generation 时:
- Worker 0 的 seen_generation = 10（第一次）
- 第二次 generation 变为 11
- Worker 0 醒来发现 seen_generation ≠ generation → 正确拿新任务
- Worker 1 醒来同样检查 generation → 正确

关键: generation 让每个 worker 能区分"我已经处理过的任务"和"新任务"

今日学习检查清单 ​

  • 能解释 mutex 和 condvar 的作用
  • 理解线程池的 worker 主循环
  • 能计算行范围分配
  • 理解 generation 计数器的作用
  • 能解释为什么 cond_wait 必须在 while 中
  • 理解 TLS 防止嵌套并行的机制

延伸挑战 ​

挑战 1（中级）：线程扩展性实验 ​

修改 ds4.c 的线程数（通过环境变量或修改 n_workers），分别用 1、2、4、8 线程运行推理。记录 decode 速度（t/s），画出"线程数 vs 速度"曲线。是否线性增长？瓶颈在哪里？

挑战 2（高级）：对比线程池和 work-stealing ​

ds4.c 的线程池用固定行范围分配（rows_per_thread = ceil(n / n_threads)）。这可能导致负载不均——如果某些行的计算量更大（如 MoE 专家命中不均）。设计一个简单的 work-stealing 方案：worker 完成自己的任务后如何"偷"其他 worker 的任务？需要哪些同步原语？