Appearance
Part 4: 代码走读
1. RoPE 与推理循环
追踪 1：YaRN 频率缩放
频率修正斜坡函数
c
// 行 4675: 线性衰减斜坡，用于平滑过渡维度区间
static float rope_yarn_ramp(float low, float high, int i0) {
const float y = ((float)(i0 / 2) - low) / fmaxf(0.001f, high - low);
return 1.0f - fminf(1.0f, fmaxf(0.0f, y));
// ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
// 当 i0 < 2*low: 返回 1.0（完全应用原始频率）
// 当 i0 > 2*high: 返回 0.0（完全应用缩放频率）
// 中间: 线性过渡
}修正维度范围计算
c
// 行 4680: 计算需要频率修正的维度边界
static float rope_yarn_corr_dim(int n_dims, uint64_t n_ctx_orig,
float n_rot, float base) {
return (float)n_dims * logf((float)n_ctx_orig / (n_rot * 2.0f * (float)M_PI))
/ (2.0f * logf(base));
}
// 行 4684: 计算 [低频修正起点, 低频修正终点]
static void rope_yarn_corr_dims(int n_dims, uint64_t n_ctx_orig,
float freq_base, float beta_fast,
float beta_slow, float dims[2]) {
// dims[0] = rope_yarn_corr_dim(..., beta_fast) ← 快速衰减开始
// dims[1] = rope_yarn_corr_dim(..., beta_slow) ← 慢速衰减结束
}维度索引 vs 频率修正:
维度: 0 2 4 6 ... dims[0] ... dims[1] ... n_rot
频率: 高 ←──────────────────────────────────────→ 低
修正: 无修正 ramp: 1→0 全修正
↑过渡区↑追踪 2：RoPE 尾部旋转应用
逐层频率配置
c
// 行 4745: 压缩层使用不同的 RoPE 频率基数
static float layer_rope_freq_base(uint32_t il) {
return ds4_layer_compress_ratio(il) != 0 && DS4_COMPRESS_ROPE_FREQ_BASE > 0.0f
? DS4_COMPRESS_ROPE_FREQ_BASE // 压缩层: 独立频率
: DS4_ROPE_FREQ_BASE; // 普通层: 全局频率
}
// 行 4751: 压缩层的频率缩放
static float layer_rope_freq_scale(uint32_t il) {
if (ds4_layer_compress_ratio(il) == 0 || DS4_ROPE_SCALE_FACTOR <= 0.0f)
return 1.0f;
return 1.0f / DS4_ROPE_SCALE_FACTOR;
}旋转应用（带 YaRN 扩展）
c
// 行 4758: 对 Q/KV 的尾部维度应用 RoPE
static void rope_tail_layer_inplace(
float *x, uint32_t n_head, uint32_t head_dim,
uint32_t n_rot, uint32_t pos, uint32_t il, bool inverse) {
const bool compressed = ds4_layer_compress_ratio(il) != 0;
const float freq_base = layer_rope_freq_base(il);
const float freq_scale = layer_rope_freq_scale(il);
const float ext_factor = compressed && DS4_ROPE_SCALE_FACTOR > 1.0f
? 1.0f : 0.0f;
// 对每个头的尾部长度应用旋转
for (uint32_t h = 0; h < n_head; h++) {
float *tail = x + (uint64_t)h * head_dim; // 尾部维度
float theta_scale = powf(freq_base, -2.0f / (float)n_rot);
for (uint32_t i = 0; i < n_rot; i += 2) {
float theta = pos * powf(theta_scale, i / 2); // θ = pos × base^(-2i/d)
// YaRN 修正: 根据 ext_factor 和 ramp 混合原始/缩放频率
const float c = cosf(theta) * mscale;
const float s = sin_sign * sinf(theta) * mscale;
// 2D 旋转
float x0 = tail[i + 0], x1 = tail[i + 1];
tail[i + 0] = x0 * c - x1 * s;
tail[i + 1] = x0 * s + x1 * c;
}
}
}RoPE 旋转矩阵:
[x0'] [cos θ -sin θ] [x0]
[x1'] = [sin θ cos θ] [x1]
θ = pos × base^(-2i/d)
pos=0: θ=0, 旋转 0° → 不旋转（位置 0 无位置信息）
pos=1: θ=θ₁, 旋转 θ₁ → 低维度旋转角度大（高频）
pos=N: θ=θₙ, 旋转 θₙ → 高维度旋转角度小（低频）
每个 token 的位置信息编码在旋转角度中，
不同维度的旋转频率不同 → 自然而然地捕获相对位置关系追踪 3：Prefill 层主序前向
批量层主序
c
// 行 7674: 所有 prompt token 依次通过每一层
static void prefill_layer_major_cpu(
float *logits, const ds4_model *model,
const ds4_weights *weights,
ds4_kv_cache *cache, const token_vec *prompt, ...) {
const uint32_t n_tok = prompt->len;
// 分配双缓冲
float *cur = xmalloc(n_tok * hc_dim * sizeof(float));
float *next = xmalloc(n_tok * hc_dim * sizeof(float));
float *attn = xmalloc(n_tok * hc_dim * sizeof(float));
// Token 嵌入
for (uint32_t t = 0; t < n_tok; t++)
embed_token_f16(cur + t * hc_dim, model, weights, prompt->data[t]);
// 层主序: 所有 token 通过层 0，再通过层 1，...
for (uint32_t il = 0; il < DS4_N_LAYER; il++) {
// 注意力子层（token 间并行）
for (uint32_t t = 0; t < n_tok; t++) {
layer_attention_prefill_one(attn + t*hc_dim, model,
layer, cache, cur, il, t, ...);
}
// FFN 子层（批量 128 token 一组）
const uint32_t ffn_batch = 128;
for (uint32_t b = 0; b < n_tok; b += ffn_batch) {
uint32_t end = MIN(b + ffn_batch, n_tok);
layer_ffn_batch(next + b*hc_dim, model, layer,
attn + b*hc_dim, prompt->data + b, end - b, ...);
}
// 交换缓冲区
float *tmp = cur; cur = next; next = tmp;
}
// 输出 logits
for (uint32_t t = 0; t < n_tok; t++)
output_logits_one(logits, model, weights, cur + t * hc_dim);
}Prefill 层主序 vs Token 主序:
层主序（ds4 使用）:
Layer 0: [tok0, tok1, ..., tokN] → 全部处理
Layer 1: [tok0, tok1, ..., tokN] → 全部处理
...
Layer 42: [tok0, tok1, ..., tokN]
Token 主序:
tok0: [Layer 0, Layer 1, ..., Layer 42] → 一个 token 走完全程
tok1: [Layer 0, Layer 1, ..., Layer 42]
...
层主序优势:
- 每层只需加载一次权重（所有 token 共享）
- FFN 可以批量处理（MoE 按专家分组）
- 内存访问模式更友好（权重驻留在缓存中）追踪 4：Decode 单 token 生成
完整层前向
c
// 行 7438: 单 token 的完整 Transformer 层
static void layer_forward_raw_swa_one(
float *out_hc, const ds4_model *model,
const ds4_layer_weights *layer,
ds4_layer_cache *cache,
const float *inp_hc, uint32_t il,
uint32_t pos, int token,
const float *steering_dirs,
float steering_attn_scale, float steering_ffn_scale,
ds4_cpu_decode_scratch *scratch) {
// 注意力子层:
// HC pre → RMSNorm → Q/KV 投影 → RoPE → KV 缓存写入
// → 注意力计算 → 输出投影 → HC post
// FFN 子层:
// HC pre → RMSNorm → routed MoE + shared FFN → HC post
// → 可选 steering
}Decode 循环
c
// 行 7649: 单 token decode 入口
static void forward_token_raw_swa_cpu(
float *logits, const ds4_model *model,
const ds4_weights *weights,
ds4_kv_cache *cache, int token, uint32_t pos) {
// 初始化暂存区（根据上下文长度估算大小）
// 逐层调用 layer_forward_raw_swa_one_decode_scratch
// 最终输出 logits
}Decode 循环:
for each token to generate:
│
▼ forward_token_raw_swa_cpu(logits, model, cache, token, pos)
│
├── Layer 0: attention(1 token, n_kv rows) + FFN(1 token)
├── Layer 1: attention(1 token, n_kv rows) + FFN(1 token)
├── ...
└── Layer 42: attention(1 token, n_kv rows) + FFN(1 token)
│
▼ output_logits_one → logits[129280]
│
▼ sample_top_p_min_p → token_id
│
pos++
继续...追踪 5：输出 Logits 与首 token
首 token 特殊处理
c
// 行 7836: 首个 token 的层前向（无 KV 缓存历史）
static void layer_forward_self_one(
float *out_hc, const ds4_model *model,
const ds4_layer_weights *layer,
const float *inp_hc, uint32_t il,
uint32_t pos, int token) {
// Pipeline: HC pre → RMSNorm → Q/KV → RoPE → FP8 量化
// → F16 舍入 → 自注意力（只有自身）→ 逆 RoPE → 输出投影
// → HC post → FFN
//
// 特殊: 无 KV 缓存写入，无压缩，无 steering
// 这是诊断模式，用于理解单 token 的行为
}
// 行 7891: 从 BOS token 开始的首 token 生成
static void forward_first_token_cpu(
float *out_hc, const ds4_model *model,
const ds4_weights *weights, int token) {
// 嵌入 token
float *cur = xmalloc(hc_dim * sizeof(float));
float *next = xmalloc(hc_dim * sizeof(float));
embed_token_f16(cur, model, weights, token);
hc_from_plain_embedding(cur, ...); // 转换为 HC 状态
// 逐层前向
for (uint32_t il = 0; il < DS4_N_LAYER; il++) {
layer_forward_self_one(next, model, layer, cur, il, 0, token);
float *tmp = cur; cur = next; next = tmp; // 双缓冲交换
}
memcpy(out_hc, cur, hc_dim * sizeof(float));
free(cur); free(next);
}完整推理流程:
用户输入: "Hello world"
│
▼ 分词 → [BOS, Hello, ▁world]
│
▼ Prefill (层主序)
│ 所有 3 个 token 依次通过 43 层
│ 每层: 注意力 + FFN + KV 缓存写入
│
▼ 输出 logits → 采样 → token_4 = "!"
│
▼ Decode 循环
│ token_4 → 43 层 → logits → 采样 → token_5 = " How"
│ token_5 → 43 层 → logits → 采样 → token_6 = " are"
│ ...
│ token_N → EOS → 停止
│
▼ 解码 → "Hello world! How are you?"2. 线程池
追踪 1：线程池结构体设计
完整字段解析
c
// 行 612: 工作函数签名
typedef void (*ds4_parallel_fn)(void *ctx, uint64_t row0, uint64_t row1);
// 行 614: 最大线程数
#define DS4_MAX_THREADS 32
// 行 616-630: 线程池结构体
typedef struct {
pthread_t threads[DS4_MAX_THREADS]; // 线程句柄数组
pthread_mutex_t mutex; // 保护以下共享字段的锁
pthread_cond_t work_cond; // 工作线程等待条件
pthread_cond_t done_cond; // 主线程等待条件
uint32_t n_threads; // 总线程数（含主线程）
uint32_t n_workers; // 工作线程数（= n_threads - 1）
uint32_t generation; // 任务代号（避免虚假唤醒）
uint32_t done; // 已完成的工作线程数
bool initialized; // 是否已初始化
bool shutdown; // 是否正在关闭
ds4_parallel_fn fn; // 当前工作函数
void *ctx; // 工作函数上下文
uint64_t n_rows; // 总行数
} ds4_thread_pool;
// 行 632-634: 全局变量
static ds4_thread_pool g_pool; // 单例线程池
static __thread int g_parallel_depth; // 线程局部递归深度
static uint32_t g_requested_threads; // 用户请求的线程数内存布局
g_pool (全局单例)
┌──────────────────────────────────────┐
│ threads[0..31] (pthread_t × 32) │ ← 线程句柄
│ mutex (pthread_mutex_t) │ ← 互斥锁
│ work_cond (pthread_cond_t) │ ← 唤醒工人
│ done_cond (pthread_cond_t) │ ← 唤醒主线程
│ n_threads = 9 │ ← 例如 8 核 + 主线程
│ n_workers = 8 │
│ generation = 42 │ ← 当前任务编号
│ done = 0 │ ← 已完成计数
│ fn = matvec_q8_0_rows_worker │ ← 当前任务函数
│ ctx = &matvec_ctx │ ← 任务上下文指针
│ n_rows = 32768 │ ← 总行数
└──────────────────────────────────────┘
线程 0 (主线程) 线程 1..8 (工作线程)
g_parallel_depth=0 g_parallel_depth=1
(独立副本) (独立副本)追踪 2：Worker 主循环与 barrier 同步
代号机制
c
// 行 636: 工作线程入口
static void *ds4_worker_main(void *arg) {
const int tid = (int)(uintptr_t)arg; // 线程编号
uint32_t seen_generation = 0; // 已处理的代号
for (;;) {
// 等待新任务
pthread_mutex_lock(&g_pool.mutex);
while (seen_generation == g_pool.generation && !g_pool.shutdown) {
pthread_cond_wait(&g_pool.work_cond, &g_pool.mutex);
// ^^^^^^^^^^^^ ^^^^^^^^^^^^
// 等待被唤醒 释放锁并等待
}
if (g_pool.shutdown) {
pthread_mutex_unlock(&g_pool.mutex);
return NULL; // 关闭退出
}
// 取任务参数（在锁内，安全）
const ds4_parallel_fn fn = g_pool.fn;
void *ctx = g_pool.ctx;
const uint64_t n_rows = g_pool.n_rows;
seen_generation = g_pool.generation;
pthread_mutex_unlock(&g_pool.mutex);
// 计算行范围
const uint64_t rows_per_thread = (n_rows + g_pool.n_threads - 1)
/ g_pool.n_threads;
const uint64_t row0 = tid * rows_per_thread;
const uint64_t row1 = MIN(row0 + rows_per_thread, n_rows);
// 执行工作（标记递归深度）
g_parallel_depth++;
fn(ctx, row0, row1); // ← 实际工作
g_parallel_depth--;
// 完成通知
__sync_fetch_and_add(&g_pool.done, 1); // 原子递增
if (g_pool.done == g_pool.n_workers) {
pthread_mutex_lock(&g_pool.mutex);
pthread_cond_signal(&g_pool.done_cond); // 唤醒主线程
pthread_mutex_unlock(&g_pool.mutex);
}
}
}时序图（generation = 42, n_workers = 3）:
主线程 Worker 0 Worker 1 Worker 2
│ │ │ │
├─ set fn, ctx ──────┤ │ │
├─ generation++ ─────┤ │ │
├─ broadcast ───────→│ ← cond_wait │ ← cond_wait │ ← cond_wait
│ │ │ │
├─ fn(ctx, 0, r/3) ├─ fn(ctx,r/3,2r/3)├─ fn(ctx,2r/3,r) │
│ 主线程执行切片0 │ 切片1 │ 切片2 │
│ │ │ │
│ ├─ done=1 ──────────┤ │
│ │ ├─ done=2 ──────────┤
│ │ │ ├─ done=3
│ ←─────────────────────────────────────── signal done_cond ─┤
│ │ │ │
▼ ▼ ▼ ▼追踪 3：线程池初始化与关闭
初始化
c
// 行 678: 初始化线程池
static void ds4_threads_init(void) {
if (g_pool.initialized) return;
// 默认: min(CPU 核心数, 12)
long nprocs = sysconf(_SC_NPROCESSORS_ONLN);
uint32_t n = (uint32_t)nprocs;
if (n > 12) n = 12; // 上限 12 线程
// 环境变量覆盖
const char *env = getenv("DS4_THREADS");
if (env) n = atoi(env);
if (n < 1) n = 1;
if (n > DS4_MAX_THREADS) n = DS4_MAX_THREADS;
g_pool.n_threads = n;
g_pool.n_workers = n - 1; // 主线程也算一个
// 创建 n-1 个工作线程（主线程是 worker 0）
for (uint32_t i = 1; i < n; i++) {
pthread_create(&g_pool.threads[i], NULL,
ds4_worker_main, (void *)(uintptr_t)i);
}
g_pool.initialized = true;
}关闭
c
// 行 715: 优雅关闭
static void ds4_threads_shutdown(void) {
if (!g_pool.initialized) return;
pthread_mutex_lock(&g_pool.mutex);
g_pool.shutdown = true; // 设置关闭标志
g_pool.generation++; // 代号递增触发唤醒
pthread_cond_broadcast(&g_pool.work_cond); // 广播唤醒所有工人
pthread_mutex_unlock(&g_pool.mutex);
// 等待所有工作线程退出
for (uint32_t i = 1; i < g_pool.n_threads; i++) {
pthread_join(g_pool.threads[i], NULL);
}
g_pool.initialized = false;
}追踪 4：parallel_for 分发机制
行范围分配策略
c
// 行 736: 带最小行数阈值的并行分发
static void ds4_parallel_for_min_rows(
uint64_t n_rows, ds4_parallel_fn fn, void *ctx,
uint64_t min_parallel_rows) {
// 退化为串行的情况:
if (g_parallel_depth > 0) goto serial; // 嵌套并行 → 串行
if (g_pool.n_threads <= 1) goto serial; // 单线程
if (n_rows < min_parallel_rows) goto serial; // 行数太少
// 并行分发
ds4_threads_init(); // 懒初始化
pthread_mutex_lock(&g_pool.mutex);
g_pool.fn = fn;
g_pool.ctx = ctx;
g_pool.n_rows = n_rows;
g_pool.done = 0;
g_pool.generation++; // 新代号
pthread_cond_broadcast(&g_pool.work_cond); // 唤醒工人
pthread_mutex_unlock(&g_pool.mutex);
// 主线程执行切片 0
{
const uint64_t rpt = (n_rows + g_pool.n_threads - 1) / g_pool.n_threads;
g_parallel_depth++;
fn(ctx, 0, MIN(rpt, n_rows)); // 主线程的行范围
g_parallel_depth--;
}
// 等待所有工人完成
pthread_mutex_lock(&g_pool.mutex);
while (g_pool.done < g_pool.n_workers) {
pthread_cond_wait(&g_pool.done_cond, &g_pool.mutex);
}
pthread_mutex_unlock(&g_pool.mutex);
return;
serial:
fn(ctx, 0, n_rows); // 直接串行执行
}
// 行 770: 默认最小行数 512 的便捷封装
static void ds4_parallel_for(uint64_t n_rows,
ds4_parallel_fn fn, void *ctx) {
ds4_parallel_for_min_rows(n_rows, fn, ctx, 512);
}行分配示例 (n_rows=32768, n_threads=9):
rows_per_thread = ceil(32768 / 9) = 3641
Thread 0 (主): [ 0, 3641)
Thread 1: [ 3641, 7282)
Thread 2: [ 7282, 10923)
...
Thread 8: [29128, 32768)
每个线程处理约 3641 行，负载基本均衡
最后一线程可能略少（32768 - 29128 = 3640）追踪 5：TLS 与并发安全
递归并行防护
c
// 行 633: 线程局部存储
static __thread int g_parallel_depth;为什么需要防止递归并行？
场景: matvec_q8_0_rows 被 parallel_for 调用
内部某个函数又调用 parallel_for
ds4_parallel_for → fn(ctx, 0, rows)
└── some_matvec_worker(...)
└── ds4_parallel_for(n, sub_fn, sub_ctx)
└── 如果不加防护 → 死锁！
所有工作线程都在执行第一个任务
没有线程空闲来执行第二个任务
防护机制:
g_parallel_depth++ → 值变为 1
ds4_parallel_for_min_rows 检测到 depth > 0
→ 退化为串行执行 fn(ctx, 0, n_rows)
g_parallel_depth-- → 恢复为 0原子操作
c
// Worker 完成通知使用原子操作（无锁）
__sync_fetch_and_add(&g_pool.done, 1); // 原子递增
if (g_pool.done == g_pool.n_workers) { // 最后一个完成
pthread_cond_signal(&g_pool.done_cond);
}done 的递增不需要 mutex 保护，因为 __sync_fetch_and_add 是编译器内置的原子操作。只有当最后一个 worker 完成时才需要锁来 signal 条件变量。