Skip to content

Part 4: 推理流程

RoPE、prefill/decode、线程池

1. RoPE 与推理循环

"我爱猫"和"猫爱我"的注意力分数一样——模型分不清词序。RoPE 通过旋转向量编码位置信息。今天还学习 prefill vs decode 两个推理阶段,以及自回归生成的本质:每步输出成为下步输入,无法并行。

RoPE 在 Self-Attention 的 Q/K 投影中注入位置信息,使模型感知 token 顺序。

设计决策推导:位置编码与推理阶段

问题 1:注意力分数只看向量内容,不知道 token 顺序
     → "我爱猫"和"猫爱我"的 Q/K 值完全相同,模型分不清
  └─ 方案:RoPE — 按 token 位置旋转 Q 和 K
     2D 旋转矩阵有性质 R(a)·R(b) = R(b-a)
     → Q·K 点积只依赖相对距离,不依赖绝对位置

问题 2:训练时只见过 64K 长度,推理时要处理 1M token?
  └─ 方案:YaRN — 混合外推和内插
     高频分量(精确区分相邻位置)→ 内插(缩小频率)
     低频分量(感知远距离关系)→ 外推(直接延伸)
     → 支持 16× 上下文扩展(64K → 1M)

问题 3:用户发来 5000 token 的 prompt,怎么处理?
  └─ Prefill 阶段 — 所有 prompt token 批量通过 43 层
     → 可以并行(矩阵乘法天然并行)
     → GPU 利用率高,速度 250+ t/s

问题 4:生成回复时,每步只产生 1 个 token,GPU 利用率极低?
  └─ Decode 阶段 — 逐 token 串行生成(自回归)
     每步要读取 81GB 权重但只算 1 个 token → 带宽瓶颈
     → 速度只有 5-10 t/s
     → 这就是 decode 比 prefill 慢 25-50 倍的根本原因

问题 5:decode 能不能加速?
  └─ 方案:MTP 推测解码 — 用小模型猜 N 个 token,
     大模型一次验证,猜对就"免费"获得 N 个 token
     → 接受率 70% 时理论加速 2×+

C 知识点

1. 三角函数

RoPE 用 cos/sin 实现旋转变换:

c
const float c = cosf(theta) * mscale;
const float s = sin_sign * sinf(theta) * mscale;
const float x0 = tail[i + 0];
const float x1 = tail[i + 1];

tail[i + 0] = x0 * c - x1 * s;    // 旋转后的第一个分量
tail[i + 1] = x0 * s + x1 * c;    // 旋转后的第二个分量

这是 2D 旋转矩阵的数值实现:

[x0']   [cos θ  -sin θ] [x0]
[x1'] = [sin θ   cos θ] [x1]

2. 幂运算与频率

c
const float theta_scale = powf(freq_base, -2.0f / (float)n_rot);
// ...
theta_extrap *= theta_scale;   // 每对维度乘以一个衰减因子

RoPE 的频率按指数衰减:第 i 对维度使用 freq_base^(-2i/n_rot) 作为频率。低维度变化快(高频),高维度变化慢(低频)。

3. 推理循环的状态管理

生成循环维护一个位置计数器 pos

c
for (int i = 0; i < n_predict && pos < ctx_size; i++) {
    int token = sample_argmax(logits, DS4_N_VOCAB);   // 选 token
    if (token == vocab->eos_id) break;                  // 遇到 EOS 停止

    if (emit) emit(emit_ud, token);                     // 回调输出
    n_generated++;

    forward_token(logits, model, weights, &cache,       // 前向传播一个 token
                  token, (uint32_t)pos, &scratch);
    pos++;                                                // 位置 +1
}

LLM 知识点

1. RoPE(Rotary Position Embedding)

RoPE 通过旋转向量来编码位置信息:

位置 0: 向量旋转 0°
位置 1: 向量旋转 θ°
位置 2: 向量旋转 2θ°
位置 n: 向量旋转 nθ°

关键性质:两个位置的内积只依赖于它们的相对距离:

Q_pos_m · K_pos_n = f(m - n)

这使得注意力自然地编码相对位置——不需要额外的位置参数。

2. YaRN 扩展

RoPE 在超出训练长度后效果下降。YaRN(Yet another RoPE extensioN)通过插值解决:

c
// 在训练长度内:正常 RoPE
// 超出训练长度:混合外推和内插
const float ramp_mix = rope_yarn_ramp(corr_dims[0], corr_dims[1], i) * ext_factor;
theta = theta_interp * (1.0f - ramp_mix) + theta_extrap * ramp_mix;

低维度(高频):保持外推(因为它们已经学到了长距离关系) 高维度(低频):改为内插(缩放到更小的角度)

ds4.c 的配置:

  • DS4_ROPE_FREQ_BASE = 10000.0(标准 base)
  • DS4_ROPE_SCALE_FACTOR = 16.0(扩展 16 倍)
  • DS4_ROPE_ORIG_CTX = 65536(训练时上下文长度)
  • 扩展后支持 1M token 上下文

3. Prefill vs Decode

推理分两个阶段:

Prefill(预填充): 处理整个 prompt

输入: [token_0, token_1, ..., token_n]  (整个 prompt)
处理: 所有 token 同时通过所有层
输出: 最后一个 token 的 logits
复杂度: O(n × L × d²)  (n = prompt 长度)

ds4.c 用层主序(layer-major)处理:

c
// 先把所有 token 通过第 0 层,再通过第 1 层...
for (uint32_t il = 0; il < DS4_N_LAYER; il++) {
    // 批量注意力:所有 token 互相关注
    layer_attention_raw_swa_batch(attn, ..., n_tok, il, 0);
    // 批量 FFN:可以分组处理
    layer_ffn_batch(next, ..., n_tok, il);
}

Decode(解码): 逐个生成 token

输入: token_t  (一个新 token)
处理: 通过所有层,用 KV cache 做注意力
输出: 下一个 token 的 logits
复杂度: O(L × d²)  (常量,与历史长度无关... 不完全对,因为有压缩 KV)
c
// 每个 token 独立通过所有层
for (uint32_t il = 0; il < DS4_N_LAYER; il++) {
    layer_forward_raw_swa_one(next, model, &weights->layer[il],
                              &cache->layer[il], cur, il, pos, token, scratch);
    swap(cur, next);
}

性能对比:

  • Prefill: 快速,因为可以批量处理(GPU 并行)
  • Decode: 较慢,因为每次只处理一个 token(串行)

4. 自回归生成

LLM 的生成是自回归的——每一步的输出成为下一步的输入:

这个过程本质上是串行的——无法并行生成多个 token(因为每一步依赖上一步的结果)。这就是为什么 decode 速度(t/s)是 LLM 推理的关键指标。

5. 推测解码(Speculative Decoding / MTP)

ds4.c 支持 MTP(Multi-Token Prediction)推测解码来加速:

1. 主模型生成 token A
2. MTP 小模型快速推测: A → B → C → ...(多步草稿)
3. 主模型验证: 确认哪些草稿 token 正确
4. 接受正确的前缀,丢弃错误的
5. 重复

理想情况: 一次验证接受多个 token → 加速
实际情况: 接受率取决于 MTP 模型质量

ds4.c 的实现(行 16071-16320):

  • 草稿深度由 --mtp-draft N 控制;深度 2 时走专用的 argmax 快速路径
  • 有置信度门槛(mtp_margin):低置信度时跳过 MTP
  • Metal 路径上用 decode2_exact 等内核一次验证多个位置

推测验证的"逐行 argmax"不变量

推测解码的验证步必须保证一个不变量:每个被提交(commit)的草稿 token,必须是其位置上主模型 logits 的(近似)argmax——否则就把模型从没生成过的 token 当成它生成的,污染上下文。

这条不变量一度在 --mtp-draft > 2 时被破坏:metal_graph_verify_suffix_tops() 本该对 top_rows 个 logits 行各取 top-1,却误调成了"对单行取 top-top_rows"(n_tokens=1, top_k=top_rows 而非 n_tokens=top_rows, top_k=1)。于是 row_tops[i>0] 拿到的是第 0 行的亚军 token,验证器拿错误的行去匹配草稿,接受了模型根本没生成的 token。

为什么深度 ≤2 看不到这个 bug?因为 top_rows == 1 时代码走的是上方的专用 argmax 分支,根本到不了那段错误调用。所以这是一个"只在开深推测时才暴露"的隐藏缺陷:

c
// ds4_gpu_indexer_topk_tensor() 的签名是 (selected, scores, n_comp, n_tokens, top_k)
//  错误: (..., n_tokens=1,        top_k=top_rows)  // 单行的前 top_rows 名
//  正确: (..., n_tokens=top_rows, top_k=1)         // 每行的 top-1

修复后每个被提交 token 的 argmax 间距(worst gap)从 ~20.96 logits 降到 0.00。该调用点在后端共享的 ds4.c 里,所以 CUDA 同样受影响、同样修复。回归测试 tests/ds4_test.c --mtp-verify-depth:在深度 4 跑贪心推测解码一个逐字复制任务,再把提交的 token 教师强制喂回普通解码,断言每个都是其位置的近似 argmax——这恰是推测验证必须保持的不变量。

6. 完整推理流程图

2. 线程池

43 层的矩阵运算、MoE 的 256 个专家——单线程太慢。线程池把行范围分配给多个 CPU 核心并行计算,推理速度随核心数线性增长。今天学习 pthread API、worker 主循环和防嵌套并行的 TLS 技巧。

线程池为 HTTP 服务器 提供并发请求处理能力。

C 知识点

1. pthread 基础

POSIX 线程(pthread)是 C 语言的标准线程 API:

c
pthread_t thread;
pthread_create(&thread, NULL, worker_fn, arg);   // 创建线程
pthread_join(thread, NULL);                       // 等待结束
pthread_mutex_t mu = PTHREAD_MUTEX_INITIALIZER;  // 互斥锁
pthread_cond_t cv = PTHREAD_COND_INITIALIZER;    // 条件变量

2. Mutex(互斥锁)

保护共享数据的访问:

c
pthread_mutex_lock(&g_pool.mutex);     // 加锁
// ... 访问共享数据 ...
pthread_mutex_unlock(&g_pool.mutex);   // 解锁

同一时间只有一个线程可以持有锁。其他线程在 lock 处阻塞等待。

3. Condition Variable(条件变量)

用于线程间通知:

c
// 等待方(worker):
pthread_mutex_lock(&mu);
while (条件不满足)
    pthread_cond_wait(&cv, &mu);    // 释放锁 + 等待信号 + 重新获取锁
pthread_mutex_unlock(&mu);

// 通知方(主线程):
pthread_mutex_lock(&mu);
修改条件;
pthread_cond_broadcast(&cv);         // 唤醒所有等待者
pthread_mutex_unlock(&mu);

wait 必须在 while 循环中(不是 if),因为可能被虚假唤醒。

4. 线程池设计

ds4.c 的线程池(行 594-752):

主线程                          Worker 线程 1..N-1
   │                                │
   ├─ 设置 fn, ctx, n_rows          │ (阻塞在 cond_wait)
   ├─ generation++                  │
   ├─ cond_broadcast(work) ─────────┼─→ 被唤醒
   ├─ 自己也执行一部分工作           ├─ 检查 generation
   │  fn(ctx, 0, rows_per_thread)    ├─ 执行 fn(ctx, my_start, my_end)
   │                                │
   ├─ cond_wait(done) ←─────────────┼─ done++ → cond_signal(done)
   │                                │
   └─ 所有 worker 完成,继续          └─ 回到 cond_wait

关键设计:

  • 主线程也工作:不浪费一个核
  • generation 计数器:防止旧的 broadcast 唤醒新的任务
  • 行范围分配rows_per_thread = (n_rows + n_threads - 1) / n_threads
  • __thread g_parallel_depth:防止嵌套并行(递归时退化为串行)

5. TLS(Thread-Local Storage)

c
static __thread int g_parallel_depth;

每个线程有自己的副本。在 ds4_parallel_for 中:

c
g_parallel_depth++;     // 标记"当前线程已在并行执行"
fn(ctx, 0, main_row1);  // 如果 fn 内部又调用 parallel_for
g_parallel_depth--;     // → g_parallel_depth > 0 → 退化为串行

LLM 知识点

推理并行性

ds4.c 在 prefill 阶段利用多线程加速:

c
// 并行处理多个 token 的 SwiGLU
ds4_parallel_for(n_tok, swiglu_batch_worker, &ctx);

// 并行处理多个 token 的 FFN
ds4_parallel_for(n_tok, ffn_worker, &ctx);

但 decode 阶段无法并行(每次只处理 1 个 token)。多线程只在 prefill 的批量操作中有效。

线程池的最小行数阈值(512)确保小任务不会因线程同步开销而变慢。

困惑度评分(Perplexity Scoring)

--perplexity-file 模式让 ds4 评估一段文本与模型预期的吻合程度:

  1. 将文件内容分词后送入模型
  2. 跳过前 32 个 token 作为"种子"(避免 BOS 附近的异常值)
  3. 对每个后续 token 累加负对数似然(NLL)
  4. 最终报告平均 NLL 和 perplexity(exp(mean_nll)

困惑度越低,模型对这段文本越"不意外"——用于评估量化质量、比较模型变体、或验证模型行为是否正确。CLI 通过 ds4_engine_options.perplexity_file 字段传入文件路径。

GPU 功耗节流(Power Throttling)

推理引擎支持 power_percent(1-100)参数控制 GPU 使用率。低于 100 时,引擎在每个 prefill 层和 decode token 后注入 sleep:

实际 sleep 微秒 = (100 - power_percent) / power_percent × step_time

ds4_engine_set_power() API 允许运行时动态调整(agent 模式下根据用户交互状态自动调节)。设为 100 时不插入任何 sleep,满速运行。