ds4.c 学习笔记

Appearance

Part 3: 代码走读 ​

1. 公共 API 与权重结构 ​

追踪 1：权重结构层次 ​

三级结构 ​

ds4_weights                          全局模型权重
  ├── token_embd                     词嵌入矩阵 [DS4_N_EMBD, DS4_N_VOCAB]
  ├── output_hc_base/fn/scale        超连接输出头（3 个张量）
  ├── output_norm                    输出层 RMSNorm 权重
  ├── output                         最终输出投影
  └── layer[43]                      43 层，每层一个 ds4_layer_weights

ds4_layer_weights                    单层权重（行 1987-2023）
  ├── HC 注意力控制
  │     ├── hc_attn_fn/scale/base    超连接参数（3 个张量）
  ├── 注意力机制
  │     ├── attn_norm                注意力 RMSNorm
  │     ├── attn_q_a / attn_q_a_norm / attn_q_b   两级低秩 Q 投影
  │     ├── attn_kv / attn_kv_a_norm KV 投影 + Norm
  │     ├── attn_sinks               Sink 注意力参数
  │     ├── attn_output_a / attn_output_b          两级低秩输出投影
  │     ├── attn_compressor_*        KV 压缩器（条件绑定）
  │     └── indexer_*                索引器（仅 ratio==4 层）
  ├── HC FFN 控制
  │     └── hc_ffn_fn/scale/base     FFN 超连接参数
  └── FFN / MoE
        ├── ffn_norm                 FFN RMSNorm
        ├── ffn_gate_tid2eid         Hash 路由表（仅 Hash 层）
        ├── ffn_gate_inp             路由器权重
        ├── ffn_exp_probs_b          专家概率偏置（可选）
        ├── ffn_gate/up/down_exps   256 个专家权重
        └── ffn_gate/up/down_shexp  共享专家权重

ds4_mtp_weights                      多 token 预测头（行 2035-2045）
  ├── e_proj / h_proj               嵌入/隐藏投影
  ├── enorm / hnorm / norm          归一化层
  ├── hc_head_base/fn/scale         超连接头
  └── block                         一整套 ds4_layer_weights

追踪 2：权重绑定过程 ​

调用链 ​

ds4_engine_open()                        [行 16950]
  └── weights_bind(&e->weights, &e->model)  [行 2578]
        │
        ├── memset(w, 0, sizeof(*w))           清零结构体
        │
        ├── 绑定顶层张量（按 GGUF 名称查找）
        │     required_tensor(m, "token_embd.weight")     → w->token_embd
        │     required_tensor(m, "output_hc_base.weight") → w->output_hc_base
        │     required_tensor(m, "output_norm.weight")    → w->output_norm
        │     required_tensor(m, "output.weight")         → w->output
        │
        ├── for il = 0..42:                     逐层绑定
        │     compress_ratio = ds4_layer_compress_ratio(il)
        │
        │     绑定通用注意力张量:
        │     required_tensorf("blk.%d.attn_norm.weight", il)
        │     required_tensorf("blk.%d.attn_q_a.weight", il)
        │     ...
        │
        │     if (compress_ratio != 0):         压缩层额外绑定
        │       required_tensorf("blk.%d.attn_compressor_kv.weight", il)
        │
        │     if (compress_ratio == 4):         索引器层额外绑定
        │       required_tensorf("blk.%d.indexer_proj.weight", il)
        │
        │     ffn_exp_probs_b = tensor_by_namef(...)  ← 可选！
        │
        │     if (il < DS4_N_HASH_LAYER):       Hash 路由层绑定
        │       required_tensorf("blk.%d.ffn_gate_tid2eid.weight", il)
        │
        └── weights_validate_layout(w)          [行 2285] 验证所有张量维度

条件绑定的内存布局 ​

                    Layer 0-28        Layer 29-36       Layer 37-42
compress_ratio:        0                  4                  8
                    ┌──────────┐    ┌──────────┐    ┌──────────┐
attn_norm          │    ✓     │    │    ✓     │    │    ✓     │
attn_q_a/b         │    ✓     │    │    ✓     │    │    ✓     │
compressor_kv      │          │    │    ✓     │    │    ✓     │
indexer_proj       │          │    │    ✓     │    │          │
ffn_gate_tid2eid   │    ✓     │    │          │    │          │
                    └──────────┘    └──────────┘    └──────────┘

追踪 3：Engine / Session 生命周期 ​

Engine 初始化 ​

c
// 行 16950: 引擎打开
int ds4_engine_open(ds4_engine **out, const ds4_engine_options *opt) {
    ds4_engine *e = xcalloc(1, sizeof(*e));
    e->fd = -1;                              // 文件描述哨兵
    e->backend = opt->backend;               // CPU / Metal / CUDA
    e->mtp_draft_tokens = CLAMP(opt->mtp_draft_tokens, 1, 16);
    e->mtp_margin = opt->mtp_margin < 0 ? 3.0f : opt->mtp_margin;

    // 核心加载流程
    model_open(&e->model, opt->model_path, ...);  // mmap GGUF
    vocab_load(&e->vocab, &e->model);              // 分词器
    config_validate_model(&e->model);              // 参数校验
    weights_bind(&e->weights, &e->model);          // 指针绑定

    // GPU 后端初始化
    if (backend == DS4_BACKEND_METAL) {
        ds4_gpu_init();
        ds4_gpu_set_model_map(e->model.map, e->model.size);
    }
}

Session 创建与释放 ​

ds4_session_create(&sess, engine, ctx_size)   [行 17094]
  │
  ├── CPU 路径:
  │     kv_cache_init(&s->cache, ctx_size)     分配 KV 缓存
  │     cpu_decode_scratch_init(&s->scratch)   分配解码暂存区
  │     s->logits = xmalloc(DS4_N_VOCAB * sizeof(float))  129280 × 4 字节
  │
  └── GPU 路径:
        metal_graph_alloc_raw_cap(...)         分配 Metal 计算图
        s->logits = ds4_gpu_tensor_alloc(...)

ds4_session_free(s)                            [行 17142]
  ├── kv_cache_free / scratch_free  (CPU)
  ├── ds4_gpu_graph_free            (GPU)
  ├── free(s->checkpoint.data)
  ├── free/ds4_gpu_tensor_free(s->logits)
  └── free(s)

追踪 4：Session 同步与前缀复用 ​

同步策略 ​

c
// 行 17201: 将 session 状态推进到指定 prompt
int ds4_session_sync(ds4_session *s, const ds4_tokens *prompt, ...) {
场景 1: 完全匹配（checkpoint 是 prompt 的前缀）

checkpoint: [A, B, C]              已处理 3 个 token
prompt:    [A, B, C, D, E, F]      需要 6 个 token

  → 增量计算 D, E, F
  → 如果 suffix 长度 <= 8: 逐 token decode
  → 如果 suffix 长度 > 8:  批量 prefill

场景 2: 无匹配

checkpoint: [A, B, C]
prompt:    [X, Y, Z]

  → 丢弃 checkpoint
  → 从头 prefill [X, Y, Z]

公共前缀查找 ​

c
// 行 17397: O(n) 线性比较
int ds4_session_common_prefix(ds4_session *s, const ds4_tokens *prompt) {
    if (!s->checkpoint.data) return 0;          // 无缓存，直接返回 0
    int n = MIN(s->checkpoint.len, prompt->len);
    for (int i = 0; i < n; i++) {
        if (s->checkpoint.data[i] != prompt->data[i]) return i;  // 不匹配
    }
    return n;                                   // 全部匹配
}

改写验证 ​

c
// 行 17360: 从公共前缀处改写
ds4_session_rewrite_result ds4_session_rewrite_from_common(
        ds4_session *s, const ds4_tokens *prompt, int common, ...) {
    // 验证前 common 个 token 确实一致
    for (int i = 0; i < common; i++) { ... }
    // 如果 common == checkpoint.len → 可以直接 ds4_session_sync
    // 否则检查是否需要重建 KV cache
}

追踪 5：公共采样 API ​

一行代理模式 ​

c
// 行 17405: 贪心采样
int ds4_session_argmax(ds4_session *s) {
    return sample_argmax(s->logits, DS4_N_VOCAB);
}

// 行 17424: 带参采样
int ds4_session_sample(ds4_session *s,
        float temperature, int top_k, float top_p, float min_p, uint64_t *rng) {
    return sample_top_p_min_p(s->logits, DS4_N_VOCAB,
                              temperature, top_k, top_p, min_p, rng);
}

API 调用层次 ​

用户代码
  │
  ├── ds4_session_argmax(s)
  │     └── sample_argmax(s->logits, 129280)        [行 14874]
  │           └── 线性扫描最大值
  │
  └── ds4_session_sample(s, temp, k, p, min_p, rng)
        └── sample_top_p_min_p(...)                  [行 15023]
              ├── temperature <= 0 → sample_argmax   退化为贪心
              ├── top_k <= 0 → sample_full_vocab     全词表采样
              │                    ├── 快速路径 (top_p >= 1.0): 无排序
              │                    └── 排序路径: qsort + 截断
              └── top_k > 0 → 栈上插入排序 + min-p + top-p 过滤

ds4.h 的设计模式：对外只暴露 ds4_session_argmax 和 ds4_session_sample 两个简单接口，内部自动路由到最优实现。

2. 量化与矩阵运算 ​

追踪 1：量化块结构体 ​

编译期大小验证 ​

c
// 行 157: 利用数组大小为负的 typedef 触发编译错误
#define DS4_STATIC_ASSERT(name, cond) typedef char name[(cond) ? 1 : -1]
//                                                    ^^^^   ^^^^
//                                                    条件真   条件假 → 负大小 → 编译失败

四种量化块 ​

c
// 行 132: Q2_K — 2-bit 量化，84 字节/块（256 个权重）
typedef struct {
    uint8_t  scales[QK_K / 16];   // 16 个缩放因子    = 16 字节
    uint8_t  qs[QK_K / 4];        // 256×2bit/8=64 字节 = 64 字节
    uint16_t d;                    // 块缩放            = 2 字节
    uint16_t dmin;                 // 最小值偏移        = 2 字节
} block_q2_K;                     //                     = 84 字节 ✓

// 行 139: Q4_K — 4-bit 量化，144 字节/块
typedef struct {
    uint16_t d;                    // = 2 字节
    uint16_t dmin;                 // = 2 字节
    uint8_t  scales[12];           // = 12 字节
    uint8_t  qs[QK_K / 2];        // 256×4bit/8=128 字节
} block_q4_K;                     //   = 144 字节 ✓

// 行 146: Q8_K — 8-bit 量化，292 字节/块
typedef struct {
    float   d;                     // = 4 字节
    int8_t  qs[QK_K];             // 256 字节
    int16_t bsums[QK_K / 16];     // 32 字节（用于矩阵乘分块累加）
} block_q8_K;                     //   = 292 字节 ✓

// 行 152: IQ2_XXS — 超低比特，66 字节/块
typedef struct {
    uint16_t d;                    // = 2 字节
    uint16_t qs[QK_K / 8];        // 256/8=32 个索引 × 2 字节 = 64 字节
} block_iq2_xxs;                  //   = 66 字节 ✓

每个 Q2_K 字节的权重存储 ​

qs[0] = 0b11_10_01_00    ← 4 个 2-bit 权重打包在 1 字节中
        ^^ ^^ ^^ ^^
        v3 v2 v1 v0       值 0-3，需要 scales 和 d/dmin 解码为实际浮点值

追踪 2：F16 转换与 FP8 量化 ​

IEEE 754 半精度转单精度 ​

c
// 行 1514: f16 → f32 纯软件实现（ARM NEON 有硬件加速）
static inline float f16_to_f32(uint16_t h) {
    // ARM 路径: vcvt_f32_f16 硬件指令，1 周期
    // 标量路径: 手动位操作
    const uint32_t sign = (h >> 15) & 1;     // 最高位
    const uint32_t exp  = (h >> 10) & 0x1f;  // 5 位指数
    const uint32_t mant = h & 0x3ff;         // 10 位尾数
    // 非正规数处理: exp==0 && mant!=0 → while 循环正规化
}
F16 (16 bit):               F32 (32 bit):
S EEEEE MMMMMMMMMM          S EEEEEEEE MMMMMMMMMMMMMMMMMMMMMMM
1  5bit   10bit             1   8bit          23bit

指数偏移: F16=15, F32=127
尾数位差: 23 - 10 = 13（需要左移 13 位补齐）

FP8 KV 缓存原地量化 ​

c
// 行 1635: E4M3FN 动态量化（FP8 格式）
static void dsv4_fp8_kv_quantize_row_inplace_cpu(
        float *x, uint32_t head_dim, uint32_t n_rot) {
    const uint32_t n_nope = head_dim - n_rot;  // 非 RoPE 部分
    // 每 64 个元素一块
    for (uint32_t off = 0; off < n_nope; off += 64) {
        float amax = 找到块内最大绝对值;
        amax = fmaxf(amax, 1.0e-4f);           // 防止除零

        // 动态缩放: ldexp 形式的 E4M3FN
        float scale = ldexpf(1.0f, ceilf(log2f(amax / 448.0f)));
        //                                                ^^^^
        //                                    E4M3FN 最大可表示值 = 448

        for (int i = 0; i < block_size; i++) {
            float v = x[off + i] / scale;
            v = clamp(v, -448.0f, 448.0f);
            x[off + i] = dsv4_e4m3fn_dequant_cpu(v) * scale;  // 量化后立即反量化
        }
    }
}
原始 float (32bit) → 量化 FP8 (8bit) → 反量化 float (32bit)
    3.14159              0b_0_1000_011               3.125
                         ↑    ↑    ↑
                         S  E4   M3

精度损失不可避免，但节省 4× KV 缓存内存

追踪 3：IQ2_XXS 查找表点积 ​

预计算查找表 ​

c
// 行 307: 初始化 256×128×8 的三维查找表
static void iq2xxs_signed_grid_init(void) {
    // 第一步: 构建 128 种符号模式
    for (int i = 0; i < 128; i++) {
        // ksigns_iq2xs[i] → 8 个 ±1 符号
        // kmask_iq2xs[i]  → 8 位掩码
    }
    // 第二步: 组合 256 个网格点 × 128 种符号
    for (int grid_idx = 0; grid_idx < 256; grid_idx++) {
        for (int sign_idx = 0; sign_idx < 128; sign_idx++) {
            // iq2xxs_signed_grid[grid_idx][sign_idx][8]
            // = iq2xxs_grid[grid_idx] × 符号模式[sign_idx]
        }
    }
}

点积计算（3 路硬件加速） ​

c
// 行 327: 16 元素对点积
static inline int32_t dot_iq2_pair_16(
        const int8_t *grid0,      // 第一个网格查表结果（8 个值）
        const int8_t *grid1,      // 第二个网格查表结果（8 个值）
        const int8_t *q8) {       // Q8 激活值（16 个 int8）
    // ARM NEON + dotprod:
    int32x4_t s0 = vdotq_s32(vdupq_n_s32(0), v0, q8_0);  // 硬件点积
    int32x4_t s1 = vdotq_s32(vdupq_n_s32(0), v1, q8_1);

    // ARM NEON 无 dotprod:
    int16x8_t p0 = vmull_s8(v0, q8_0);   // 8 位乘法
    acc = vpaddlq_s16(p0);                // 成对加法

    // 标量回退:
    int32_t sum = 0;
    for (int i = 0; i < 8; i++)
        sum += grid0[i] * q8[i] + grid1[i] * q8[8 + i];
}
IQ2_XXS 反量化流水线:

索引 qs[i] (8-bit) → 查表 iq2xxs_grid[qs[i]] → 8 个 2-bit 值
                        ↓
              × 符号模式 → 8 个 int8 值
                        ↓
              · Q8 激活 → 点积累加

追踪 4：Q8_0 激活量化与 matvec ​

Q8_0 量化（块大小 32） ​

c
// 行 3123: 激活值 float → int8 量化
static void quantize_q8_0_activation(
        const float *x, int8_t *xq, float *scale, uint64_t n) {
    const uint64_t blocks = (n + 31) / 32;   // 块大小 32（不是 256）

    for (uint64_t b = 0; b < blocks; b++) {
        // 找块内最大绝对值
        float amax = 0.0f;
        for (int i = 0; i < 32; i++)
            amax = fmaxf(amax, fabsf(x[b * 32 + i]));

        float d = amax / 127.0f;             // 缩放因子
        float id = (d != 0.0f) ? 1.0f / d : 0.0f;

        scale[b] = d;                        // 存储缩放因子
        for (int i = 0; i < 32; i++) {
            xq[b * 32 + i] = (int8_t)clamp(lrintf(x[b*32+i] * id), -128, 127);
        }
    }
}

Decode 热路径 matvec ​

c
// 行 3494: 单 token Q8_0 矩阵向量乘
static void matvec_q8_0(float *out, const ds4_model *m,
                         const ds4_tensor *w, const float *x) {
    matvec_q8_0_rows(out, m, w, x, 0, w->dim[1]);  // 委托给行范围版本
}

// 行 3510: Decode 热路径（零分配）
static void matvec_q8_0_decode_scratch(
        float *out, const ds4_model *m, const ds4_tensor *w,
        const float *x, ds4_cpu_decode_scratch *scratch) {
    cpu_decode_quantize_q8_0(scratch, x, in_dim);  // 量化到预分配暂存区
    matvec_q8_0_prequant(out, m, w, scratch);       // 用预量化结果做 matvec
}
Decode 每层的 matvec 调用:

float x[4096]  →  quantize_q8_0  →  int8 xq[4096] + float scale[128]
                                          │
                                          ▼ matvec_q8_0_prequant
                                          │
float out[dim]  ←  Q8_0 权重矩阵 × Q8_0 激活  ←  权重已在模型文件中预量化

追踪 5：Q2_K / IQ2_XXS 专家 matvec ​

IQ2_XXS 专家对投影 ​

c
// 行 3775: 同时计算一个专家的 gate 和 up 投影
static void matvec_iq2_xxs_expert_pair_prequant(
        float *out0, float *out1,             // gate 输出和 up 输出
        const ds4_model *m,
        const ds4_tensor *w0, *w1,            // gate 和 up 权重
        const block_q8_K *xq,                 // 共享的 Q8_K 激活
        uint32_t expert) {                    // 专家编号
    // 验证两个张量同为 IQ2_XXS 类型
    // 计算专家偏移: base + expert * expert_bytes
    // 填充 pair_ctx，调用并行内核
}

Q2_K 专家 Down 投影 ​

c
// 行 3909: 单专家 Q2_K 下投影（诊断/追踪用）
static void matvec_q2_k_expert(
        float *out, const ds4_model *m,
        const ds4_tensor *w, const float *x,
        uint32_t expert) {
    // 注: 这个版本内部做 Q8_K 量化（xmalloc），不在热路径
    block_q8_K *xq = xmalloc(...);              // 临时分配
    ds4_quantize_row_q8_K(x, xq, in_dim);       // 量化激活
    ds4_parallel_for(out_dim, matvec_q2_k_worker, &ctx);
    free(xq);                                    // 释放临时缓冲
}

MoE matvec 数据流 ​

输入激活 float x[in_dim]
    │
    ▼ quantize Q8_K（一次）
    │
    xq[QK_K 块...]  ← 所有 6 个专家共享
    │
    ├── 专家 0: matvec_iq2_xxs_expert_pair(w_gate[0], w_up[0], xq)
    │              → gate[0][mid_dim], up[0][mid_dim]
    │
    ├── 专家 1: matvec_iq2_xxs_expert_pair(w_gate[1], w_up[1], xq)
    │              → gate[1][mid_dim], up[1][mid_dim]
    │
    ├── ... (共 6 个专家)
    │
    ▼ SwiGLU: mid[i] = silu(gate[i]) * up[i]
    │
    ▼ quantize Q8_K（per-expert）
    │
    ├── 专家 0: matvec_q2_k(w_down[0], midq[0]) → out += weight[0]
    ├── 专家 1: matvec_q2_k(w_down[1], midq[1]) → out += weight[1]
    └── ...
    │
    ▼ 加权累加
    │
float out[out_dim] = Σ weight[i] * expert_down[i]

3. 注意力机制 ​

追踪 1：RMS 归一化 ​

三种 RMSNorm 变体 ​

c
// `rms_norm_no_weight()` 中（约 2700 行）: 无权重（用于 HC 控制向量）
static void rms_norm_no_weight(float *out, const float *x,
                                uint64_t n, float eps) {
    double ss = 0.0;                        // double 精度累加
    for (uint64_t i = 0; i < n; i++) ss += (double)x[i] * x[i];
    float scale = 1.0f / sqrtf((float)(ss / n) + eps);
    for (uint64_t i = 0; i < n; i++) out[i] = x[i] * scale;
}

// `rms_norm_weight()` 中（约 2709 行）: 带可学习权重（用于层归一化）
static void rms_norm_weight(float *out, const float *x,
                             const float *weight, uint64_t n, float eps) {
    // ... 同上计算 scale ...
    for (uint64_t i = 0; i < n; i++)
        out[i] = x[i] * scale * weight[i];  // 逐通道乘以学习到的 γ
}

// `head_rms_norm_inplace()` 中（约 2718 行）: 单头原地归一化（用于 Q/KV 投影后）
static void head_rms_norm_inplace(float *x, uint32_t n_head,
                                   uint32_t head_dim, float eps) {
    for (uint32_t h = 0; h < n_head; h++) {
        float *xh = x + (uint64_t)h * head_dim;
        rms_norm_no_weight(xh, xh, head_dim, eps);  // 每头独立归一化
    }
}
RMSNorm 公式:  out[i] = x[i] / sqrt(mean(x²) + ε) × weight[i]

LayerNorm vs RMSNorm:
  LayerNorm: (x - μ) / sqrt(σ² + ε) × γ + β     ← 减均值、有偏移
  RMSNorm:   x / sqrt(mean(x²) + ε) × γ           ← 不减均值、无偏移

RMSNorm 省略了均值计算和偏移参数，计算更简单，推理更快

追踪 2：HC 头压缩（MLA） ​

Sinkhorn 分解 ​

c
// `hc_split_sinkhorn_one()` 中（约 4186 行）: 将混合向量分解为 pre/post/comb 三个输出
static void hc_split_sinkhorn_one(
        float *out,           // 输出: [2*n_hc + n_hc*n_hc] 元素
        const float *mix,     // 输入: 混合系数
        const float *scale,   // [pre_scale, post_scale, comb_scale]
        const float *base,    // 基础偏置
        int n_hc, int iters, float eps) {
    // Sigmoid 生成 pre 权重: 1/(1 + exp(-z))，z = mix[i]*pre_scale + base[i]
    // 双重归一化的 comb 矩阵（Sinkhorn 迭代）
}

HC Pre / Post 操作 ​

c
// `hc_pre_from_state_one()` 中（约 4317 行）: 从隐藏状态计算 HC 控制向量
static void hc_pre_from_state_one(
        const ds4_model *model,
        const ds4_tensor *fn,            // 投影权重
        const ds4_tensor *scale_tensor,  // 缩放参数
        const ds4_tensor *base_tensor,   // 基础偏置
        const float *residual_hc,        // HC 残差状态
        float *out, float *post, float *comb) {
    // 分配暂存区，调用 hc_pre_from_state_one_scratch
    // 输出: pre 权重、post 门控、comb 组合矩阵
}

// `hc_post_one()` 中（约 4366 行）: 将子层输出与 HC 残差合并
static void hc_post_one(
        float *out_hc, const float *block_out,
        const float *residual_hc, const float *post,
        const float *comb, uint32_t n_embd, uint32_t n_hc) {
    for (uint32_t dst = 0; dst < n_hc; dst++) {
        for (uint32_t d = 0; d < n_embd; d++) {
            out_hc[dst * n_embd + d] = block_out[d] * post[dst];
            for (uint32_t src = 0; src < n_hc; src++) {
                out_hc[dst*n_embd + d] += comb[dst + src*n_hc]
                                         * residual_hc[src*n_embd + d];
            }
        }
    }
}
HC（Hyper-Connection）数据流:

residual_hc[n_hc × n_embd]
    │
    ▼ hc_pre_from_state_one
    │
    ├── pre[n_hc]      ← 每个子层输入的门控
    ├── post[n_hc]     ← 每个子层输出的门控
    └── comb[n_hc×n_hc] ← 子层间的路由矩阵
    │
    ▼ 子层计算（注意力或 FFN）
    │
    block_out[n_embd]
    │
    ▼ hc_post_one
    │
    out_hc[dst] = block_out × post[dst]
                + Σ comb[dst,src] × residual_hc[src]

追踪 3：Q / KV 投影 ​

两级低秩 Q 投影 ​

c
// `layer_q_projection_normed_one()` 中（约 4596 行）: Q = norm(proj_b(norm(proj_a(x))))
static void layer_q_projection_normed_one(
        const ds4_model *model, const ds4_layer_weights *layer,
        const float *norm, float *q) {
    // 第一级: 4096 → 1024（低秩压缩）
    float qr[1024];  // 中间维度
    matvec_q8_0(qr, model, layer->attn_q_a, norm);

    // 中间归一化
    rms_norm_weight(qr, qr, layer->attn_q_a_norm, 1024, DS4_RMS_EPS);

    // 第二级: 1024 → 64×512 = 32768（展开到所有头）
    matvec_q8_0(q, model, layer->attn_q_b, qr);

    // 头级归一化
    head_rms_norm_inplace(q, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_RMS_EPS);
}

KV 投影（单头） ​

c
// `layer_kv_projection_normed_one()` 中（约 4633 行）: KV = norm(proj(x))，输出维度 = DS4_N_HEAD_DIM = 512
static void layer_kv_projection_normed_one(
        const ds4_model *model, const ds4_layer_weights *layer,
        const float *normed, float *kv) {
    float raw[DS4_N_HEAD_DIM];  // 512
    matvec_q8_0(raw, model, layer->attn_kv, normed);
    rms_norm_weight(kv, raw, layer->attn_kv_a_norm,
                    DS4_N_HEAD_DIM, DS4_RMS_EPS);
}
DeepSeek V4 的 GQA (Grouped Query Attention):

      Q: 64 头 × 512 维 = 32768 维
     KV:  1 头 × 512 维 =   512 维

  Q 投影: 4096 → 1024 → 32768 (两级低秩)
  KV 投影: 4096 → 512 (单级)

  所有 64 个 Q 头共享 1 组 KV → GQA 64:1
  KV 体积减小 64×，是 MLA 的核心优化

追踪 4：单 token 注意力计算 ​

Sink-Aware 注意力 ​

c
// `layer_attention_rows_one()` 中（约 4897 行）: 带 Sink 的缩放点积注意力
static void layer_attention_rows_one(
        float *out_heads, const ds4_model *model,
        const ds4_layer_weights *layer,
        const float *q, const float *kv_rows, uint32_t n_kv) {

    for (uint32_t h = 0; h < DS4_N_HEAD; h++) {
        const float *qh = q + (uint64_t)h * DS4_N_HEAD_DIM;  // 当前头的 Q
        float *oh = out_heads + (uint64_t)h * DS4_N_HEAD_DIM; // 当前头输出

        // 加载 Sink（可学习的注意力汇聚点）
        const float sink_logit = tensor_data(layer->attn_sinks, ...)[h];

        float scores[512];  // 栈缓冲（n_kv <= 512 时）
        // 或堆分配（n_kv > 512 时）

        float max_score = sink_logit;   // Sink 参与最大值计算

        // 计算注意力分数
        for (uint32_t r = 0; r < n_kv; r++) {
            scores[r] = dot_f32(qh, kv_rows + r * DS4_N_HEAD_DIM,
                                DS4_N_HEAD_DIM) / sqrtf(DS4_N_HEAD_DIM);
            if (scores[r] > max_score) max_score = scores[r];
        }

        // Softmax
        float sum = expf(sink_logit - max_score);  // Sink 贡献分母但不贡献分子
        for (uint32_t r = 0; r < n_kv; r++) {
            scores[r] = expf(scores[r] - max_score);
            sum += scores[r];
        }
        float inv_sum = 1.0f / sum;

        // 加权求值（Sink 不贡献值向量）
        memset(oh, 0, DS4_N_HEAD_DIM * sizeof(float));
        for (uint32_t r = 0; r < n_kv; r++) {
            axpy_f32(scores[r] * inv_sum,
                     kv_rows + r * DS4_N_HEAD_DIM, oh, DS4_N_HEAD_DIM);
        }
    }
}
Sink-Aware Softmax:

标准:  softmax(scores) = exp(score) / Σ exp(score)
Sink:  分母 = exp(sink) + Σ exp(score)
       分子 = 仅 Σ score[r] × value[r]
       （sink 不对应任何 value 向量）

效果: sink 作为"虚拟 token"吸收多余的注意力权重，
      防止首个真实 token 被过度关注（StreamingLLM 的核心思想）

追踪 5：KV 缓存管理与混合注意力 ​

KV 缓存写入 ​

c
// `kv_cache_push_raw()` 中（约 6307 行）: 原始 SWA（滑动窗口注意力）缓存
static void kv_cache_push_raw(ds4_layer_cache *cache, const float *kv) {
    if (cache->n_raw < cache->raw_cap) {
        // 未满: 直接追加，F16 精度舍入
        float *dst = cache->raw + (uint64_t)cache->n_raw * DS4_N_HEAD_DIM;
        for (int i = 0; i < DS4_N_HEAD_DIM; i++)
            dst[i] = f16_to_f32(f32_to_f16(kv[i]));  // float → f16 → float
    } else {
        // 已满: 滑动窗口，最老的行被丢弃
        memmove(cache->raw, cache->raw + DS4_N_HEAD_DIM,
                (cache->n_raw - 1) * DS4_N_HEAD_DIM * sizeof(float));
        // 新行写到末尾
    }
    cache->n_raw++;
}

// `kv_cache_push_comp()` 中（约 6322 行）: 压缩缓存（追加式，无滑动窗口）
static void kv_cache_push_comp(float *rows, uint32_t *n_rows,
                                uint32_t cap_rows, uint32_t row_dim,
                                const float *kv) {
    // 超过容量直接报错（压缩缓存不支持丢弃）
    rows[(*n_rows)++] = kv[...];  // 同样 F16 舍入
}

混合注意力 ​

c
// `layer_attention_mixed_one()` 中（约 6613 行）: 同时在 raw 和 comp 缓存上做注意力
static void layer_attention_mixed_one(
        float *out_heads, const ds4_model *model,
        const ds4_layer_weights *layer,
        const float *q,
        const float *raw_kv, uint32_t n_raw,
        const float *comp_kv, uint32_t n_comp,
        const bool *comp_allowed) {       // 哪些 comp 行可参与
    const uint32_t n_total = n_raw + n_comp;

    // 先算 raw 行的分数
    for (uint32_t r = 0; r < n_raw; r++) { ... }

    // 再算 comp 行的分数（跳过 comp_allowed[r]==false 的行）
    for (uint32_t r = 0; r < n_comp; r++) {
        if (!comp_allowed[r]) continue;   // ratio-4 层的间歇性压缩
        ...
    }
}
混合注意力内存布局:

┌─────────────────────── Raw SWA 缓存 ──────────────────────┐
│ [tok_0] [tok_1] ... [tok_W-1]  ← 滑动窗口，最多 W 行     │
└──────────────────────────────────────────────────────────┘
                          +
┌─────────────────────── Compressed 缓存 ──────────────────┐
│ [comp_0] [comp_1] ... [comp_N]  ← 累积压缩，最多 N 行     │
└──────────────────────────────────────────────────────────┘

总注意力行数 = n_raw + n_comp
Raw 保留最近 W 个 token 的精确 KV
Comp 存储更早的 token 的压缩表示
两者在同一个 softmax 中统一计算

4. MoE 混合专家 ​

追踪 1：SiLU / Softplus / SwiGLU 激活 ​

三种激活函数 ​

c
// 行 5012: SiLU (Sigmoid Linear Unit) = x × σ(x)
static float silu(float x) {
    return x * sigmoid_stable(x);
}

// 行 5016: 数值稳定的 Softplus = log(1 + exp(x))
static float softplus_stable(float x) {
    if (x > 20.0f)  return x;           // 大正数: log(1+exp(x)) ≈ x
    if (x < -20.0f) return expf(x);     // 大负数: log(1+exp(x)) ≈ exp(x)
    return log1pf(expf(x));             // 中间区间: 精确计算
}

// 行 5022: SwiGLU = SiLU(gate) × up
static void swiglu(float *out, const float *gate,
                    const float *up, uint64_t n) {
    for (uint64_t i = 0; i < n; i++) {
        out[i] = silu(gate[i]) * up[i];
    }
}
SiLU 图形:
          |        /
    3     |      /
    2     |    /
    1     |  /
    0 ----+--------
   -1     |  \      ← x < 0 时有轻微负输出（"自门控"特性）
          |
          +---------------→ x

SwiGLU 数据流:
gate[i] → SiLU → × up[i] → out[i]
                ↑
         门控信号    信息流

追踪 2：共享 FFN（密集前馈） ​

单 token 共享专家 ​

c
// 行 5029: 共享专家的完整前馈
static void layer_shared_ffn_one(
        float *out, const ds4_model *model,
        const ds4_layer_weights *layer, const float *x) {

    // gate 和 up 投影同时计算
    float *gate = xmalloc(DS4_N_FF_EXP * sizeof(float));
    float *up   = xmalloc(DS4_N_FF_EXP * sizeof(float));
    float *mid  = xmalloc(DS4_N_FF_EXP * sizeof(float));

    // 激活量化 + pair matvec（gate 和 up 共享量化激活）
    quantize_q8_0_activation(x, xq, scale, in_dim);
    matvec_q8_0_pair_prequant(gate, up, model,
                               layer->ffn_gate_shexp,
                               layer->ffn_up_shexp, xq, scale);

    // SwiGLU 激活
    swiglu(mid, gate, up, DS4_N_FF_EXP);

    // Down 投影
    matvec_q8_0(out, model, layer->ffn_down_shexp, mid);

    free(gate); free(up); free(mid);
}

Decode 热路径（零分配） ​

c
// 行 5063: 使用预分配暂存区的版本
static void layer_shared_ffn_one_decode_scratch(
        float *out, const ds4_model *model,
        const ds4_layer_weights *layer,
        const float *x, ds4_cpu_decode_scratch *scratch) {
    // 所有中间缓冲区从 scratch 中分配，无 malloc
}

批量并行 SwiGLU ​

c
// 行 5094: 批量 worker
static void swiglu_batch_worker(void *vctx, uint64_t t0, uint64_t t1) {
    for (uint64_t t = t0; t < t1; t++)
        swiglu(mid + t * DS4_N_FF_EXP,
               gate + t * DS4_N_FF_EXP,
               up + t * DS4_N_FF_EXP, DS4_N_FF_EXP);
}

追踪 3：Hash 路由与 Top-K 选择 ​

Hash 查表路由 ​

c
// 行 5148: 直接查表获取被选专家
static void layer_hash_selected_experts(
        int selected[DS4_N_EXPERT_USED],   // 输出: 6 个专家 ID
        const ds4_model *model,
        const ds4_layer_weights *layer,
        int token) {                        // 当前 token ID
    // ffn_gate_tid2eid 是预计算的 [DS4_N_EXPERT_USED, n_vocab] 查找表
    // selected[i] = table[i * n_vocab + token]
    // 无需实时计算路由，O(1) 查表
}

路由概率计算 ​

c
// 行 5169: 基于路由器的概率
static void layer_router_probs_one(
        float probs[DS4_N_EXPERT],          // 256 个专家的概率
        const ds4_model *model,
        const ds4_layer_weights *layer,
        const float *x) {
    float logits[DS4_N_EXPERT];
    matvec_q8_0(logits, model, layer->ffn_gate_inp, x);
    for (int i = 0; i < DS4_N_EXPERT; i++)
        probs[i] = sqrtf(softplus_stable(logits[i]));
    //               ^^^^^^^               ^^^^^^^^^^^^
    //               开方使分布更尖锐        softplus 保证非负
}

插入排序 Top-K ​

c
// 行 5211: O(n×k) 的插入排序 top-k
static void topk_desc(const float *score, int n, int k, int *idx) {
    for (int i = 0; i < k; i++) idx[i] = -1;         // 初始化
    for (int i = 0; i < n; i++) {                      // 扫描 n=256 个
        for (int j = 0; j < k; j++) {                   // 与当前 top-k 比较
            if (idx[j] < 0 || score[i] > score[idx[j]]) {
                for (int m = k - 1; m > j; m--)          // 后移
                    idx[m] = idx[m - 1];
                idx[j] = i;                               // 插入
                break;
            }
        }
    }
}
Top-K 选择过程 (n=256, k=6):

score[0] = 0.3  → idx = [0, -1, -1, -1, -1, -1]
score[1] = 0.8  → idx = [1, 0, -1, -1, -1, -1]
score[2] = 0.1  → idx 不变（太小）
score[3] = 0.5  → idx = [1, 3, 0, -1, -1, -1]
...
score[255] = 0.9 → idx = [255, 1, 3, 7, 42, 0]

最终选中 6 个专家: [255, 1, 3, 7, 42, 0]

追踪 4：路由 MoE 完整流水线 ​

单 token MoE ​

c
// 行 5278: 路由 MoE 的核心
static void layer_routed_moe_one(
        float *out, const ds4_model *model,
        const ds4_layer_weights *layer,
        const float *x, uint32_t il, int token,
        float clamp, bool trace) {

    // 1. 路由选择
    if (layer->ffn_gate_tid2eid) {
        layer_hash_selected_experts(selected, model, layer, token);  // Hash 查表
        layer_hash_router_weights_one(weight, model, layer, x);      // 计算权重
    } else {
        layer_topk_selected_experts(selected, weight, model, layer, x);  // Top-K 选择
    }

    // 2. 快速路径: 所有专家的 gate+up 一起计算
    matvec_iq2_xxs_experts_mid_prequant(
        gate_all, up_all, model, layer,
        selected, weight, xq, clamp);

    // 3. SwiGLU 激活（逐专家）
    for (int e = 0; e < DS4_N_EXPERT_USED; e++)
        swiglu(mid + e * DS4_N_FF_EXP,
               gate + e * DS4_N_FF_EXP,
               up + e * DS4_N_FF_EXP, DS4_N_FF_EXP);

    // 4. 量化 mid 并做 down 投影
    ds4_quantize_row_q8_K(mid, midq, ...);
    matvec_q2_k_experts_accum_prequant(
        out, model, layer, selected, weight, midq);
}

MoE 数据流 ​

输入 x[in_dim]
    │
    ├─→ 路由选择 → selected[6], weight[6]
    │
    ├─→ Q8_K 量化 x → xq
    │
    ├── 专家 0: gate[0],up[0] = IQ2_XXS_matvec(w_gate[0],w_up[0],xq)
    ├── 专家 1: gate[1],up[1] = IQ2_XXS_matvec(w_gate[1],w_up[1],xq)
    ├── ...
    ├── 专家 5: gate[5],up[5] = IQ2_XXS_matvec(w_gate[5],w_up[5],xq)
    │
    ├── mid[0] = SwiGLU(gate[0], up[0])
    ├── mid[1] = SwiGLU(gate[1], up[1])
    ├── ...
    │
    ├── Q8_K 量化 mid → midq
    │
    ├── out += weight[0] × Q2_K_matvec(w_down[0], midq[0])
    ├── out += weight[1] × Q2_K_matvec(w_down[1], midq[1])
    ├── ...
    └── out += weight[5] × Q2_K_matvec(w_down[5], midq[5])

追踪 5：FFN 层完整前向 ​

单 token FFN（shared + routed） ​

c
// 行 5576: 完整的 FFN 层
static void layer_ffn_one(
        float *out_hc, const ds4_model *model,
        const ds4_layer_weights *layer,
        const float *inp_hc, uint32_t il, int token,
        const float *steering_dirs, float steering_scale,
        bool trace) {

    // 1. HC 前置（提取 pre/post/comb 控制向量）
    hc_pre_from_state_one(model, fn, scale, base, inp_hc,
                          pre, post, comb);

    // 2. 注意力层 RMSNorm 后的状态
    rms_norm_weight(norm, inp_hc, layer->ffn_norm, DS4_N_EMBD, DS4_RMS_EPS);

    // 3. 路由 MoE
    layer_routed_moe_one(moe, model, layer, norm, il, token, clamp, trace);

    // 4. 共享专家
    layer_shared_ffn_one(shared, model, layer, norm);

    // 5. 合并
    for (int i = 0; i < DS4_N_EMBD; i++)
        ffn_out[i] = moe[i] + shared[i];

    // 6. 可选 steering 投影
    if (steering_dirs) { ... }

    // 7. HC 后置（将输出写回 HC 状态）
    hc_post_one(out_hc, ffn_out, inp_hc, post, comb, n_embd, n_hc);
}

FFN 层全貌 ​

inp_hc[n_hc × n_embd]            ← HC 残差状态
    │
    ├─→ HC pre → pre, post, comb  ← 控制向量
    │
    ├─→ RMSNorm → norm            ← 归一化
    │
    ├──→ 路由 MoE(norm) → moe     ← 6/256 个专家
    │
    ├──→ 共享 FFN(norm) → shared  ← 1 个共享专家
    │
    ├─→ moe + shared → ffn_out    ← 合并
    │
    ├─→ [steering projection]     ← 可选方向控制
    │
    └─→ HC post → out_hc          ← 写回 HC 状态

DeepSeek V4 FFN 设计: 每层 256 个路由专家 + 1 个共享专家
                       每个 token 只激活 6 个路由专家
                       稀疏度 = 1 - 6/256 ≈ 97.7%