Part 6: 代码走读

1. GPU 后端 + 总结

追踪 1：Metal 初始化与管线编译

GPU 初始化

// 行 2661: 单例初始化 Metal 设备
int ds4_gpu_init(void) {
    g_device = MTLCreateSystemDefaultDevice();     // 获取 GPU 设备
    g_queue  = [g_device newCommandQueue];          // 命令队列

    // 加载 Metal Shading Library（编译好的 .metallib）
    // 从程序嵌入的 Metal 源码编译
    g_library = [g_device newLibraryWithSource:metal_source
                                       options:nil
                                         error:&err];

    // 初始化缓存字典
    g_model_buffer_cache  = [NSMutableDictionary new];  // 模型缓冲区
    g_pipeline_cache      = [NSMutableDictionary new];  // 管线状态
    g_transient_buffers   = [NSMutableArray new];        // 临时缓冲区
}

管线状态缓存

// 行 624: 按名称获取或创建计算管线
static id<MTLComputePipelineState> ds4_gpu_get_pipeline(
        const char *function_name) {
    NSString *key = @(function_name);

    // 查缓存
    id<MTLComputePipelineState> pipe = g_pipeline_cache[key];
    if (pipe) return pipe;                      // 命中

    // 缓存未命中: 从库中查找函数并编译管线
    id<MTLFunction> fn = [g_library newFunctionWithName:key];
    pipe = [g_device newComputePipelineStateWithFunction:fn
                                                   error:&err];
    g_pipeline_cache[key] = pipe;               // 存入缓存
    return pipe;
}

Metal 初始化层次:

MTLDevice (GPU 硬件抽象)
  └── MTLCommandQueue (命令提交队列)
        └── MTLCommandBuffer (一次提交的命令批次)
              └── MTLComputeCommandEncoder (计算命令编码器)
                    └── MTLComputePipelineState (着色器管线)
                          └── MTLFunction (从 .metal 源码编译)

管线缓存:
  "matmul_q8_0"        → pipeline_state_0
  "flash_attn_prefill" → pipeline_state_1
  "routed_moe"         → pipeline_state_2
  ...
  (~60 个推理内核)

---

追踪 2：零拷贝模型映射与张量管理

模型映射

// 行 4399: 将 mmap 的模型文件映射到 Metal 缓冲区
int ds4_gpu_set_model_map(const void *model_map, uint64_t model_size) {
    // 委托给 ds4_gpu_set_model_map_range
    // 在 macOS 上使用 MTLResourceStorageModeShared
    // GPU 和 CPU 共享同一块物理内存
    // → 零拷贝: 权重数据不需要从 CPU 复制到 GPU
}

张量分配

// 行 3767: 分配 GPU 张量
ds4_gpu_tensor *ds4_gpu_tensor_alloc(uint64_t bytes) {
    // 创建 Metal Buffer（共享内存模式）
    id<MTLBuffer> buffer = [g_device newBufferWithLength:bytes
                                                 options:MTLResourceStorageModeShared];
    // 包装为 DS4MetalTensor 对象
    DS4MetalTensor *obj = [[DS4MetalTensor alloc] init];
    obj.buffer = buffer;
    obj.offset = 0;
    obj.bytes  = bytes;

    // 追踪分配统计
    g_live_bytes  += bytes;
    g_peak_bytes   = MAX(g_peak_bytes, g_live_bytes);

    // __bridge_retained: Objective-C → C 所有权转移
    return (__bridge_retained ds4_gpu_tensor *)obj;
}

张量写入

// 行 3859: CPU → GPU 数据传输
int ds4_gpu_tensor_write(ds4_gpu_tensor *tensor, uint64_t offset,
                          const void *data, uint64_t bytes) {
    DS4MetalTensor *obj = (__bridge DS4MetalTensor *)tensor;
    // 验证边界
    if (offset + bytes > obj.bytes) return -1;
    // 直接 memcpy（共享内存模式，CPU/GPU 看到同一块内存）
    memcpy((uint8_t *)[obj.buffer contents] + obj.offset + offset,
           data, bytes);
    return 0;
}

零拷贝内存架构:

CPU 虚拟地址空间                      GPU
┌─────────────────┐
│ mmap 区域        │ ← 81GB GGUF 文件
│ (权重数据)       │
│                  │    MTLResourceStorageModeShared
│ ─ ─ ─ ─ ─ ─ ─ ─│ ← 同一块物理内存
│                  │
│ Metal Buffer     │ ← ds4_gpu_tensor_alloc 分配
│ (激活值/输出)    │
└─────────────────┘

CPU 通过指针直接读写 → GPU 通过同一物理地址访问
无需 cudaMemcpy / [buffer didModifyRange]

---

追踪 3：GPU matmul 与 matvec

Q8_0 量化矩阵乘

// 行 4919: Q8_0 权重 × Q8_0 激活的矩阵乘
int ds4_gpu_matmul_q8_0_tensor(
        ds4_gpu_tensor *out,            // 输出张量
        const void *model_map,          // 模型 mmap 基地址
        uint64_t model_size,
        uint64_t weight_offset,         // 权重在 mmap 中的偏移
        uint64_t in_dim, uint64_t out_dim,
        const ds4_gpu_tensor *x,        // 输入激活张量
        uint64_t n_tok) {              // token 数量

    // 要求 in_dim % 32 == 0（Q8_0 块大小对齐）
    // 获取 Metal 缓冲区
    id<MTLBuffer> out_buf  = ...;
    id<MTLBuffer> x_buf    = ...;
    id<MTLBuffer> w_buf    = model_map + weight_offset;  // 零拷贝权重

    // 设置计算管线和参数
    id<MTLComputePipelineState> pipe = ds4_gpu_get_pipeline("matmul_q8_0");
    id<MTLComputeCommandEncoder> enc = [cb computeCommandEncoder];
    [enc setComputePipelineState:pipe];
    [enc setBuffer:out_buf offset:0 atIndex:0];
    [enc setBuffer:w_buf  offset:0 atIndex:1];
    [enc setBuffer:x_buf  offset:0 atIndex:2];
    // ... 设置其他参数 ...

    // 分发计算线程
    [enc dispatchThreadgroups:threadgroups
           threadsPerThreadgroup:threads];
    [enc endEncoding];
}

F16 矩阵乘

// 行 5141: FP16 权重矩阵乘（用于 embedding 等未量化的权重）
int ds4_gpu_matmul_f16_tensor(
        ds4_gpu_tensor *out,
        const void *model_map, uint64_t model_size,
        uint64_t weight_offset,
        uint64_t in_dim, uint64_t out_dim,
        const ds4_gpu_tensor *x,
        uint64_t n_tok) {
    // 与 Q8_0 版本结构相同
    // 但使用 "matmul_f16" 着色器管线
    // 无 in_dim 对齐要求
}

GPU matvec 线程分发:

权重矩阵 [out_dim × in_dim]
    │
    ▼ 分割为 threadgroups
    │
    每个 threadgroup:
      处理一个输出行的子范围
      加载权重块到共享内存
      与输入向量做分块点积
      累加结果

threadgroup 大小: 例如 [64, 1, 1]
总 threadgroups: [ceil(out_dim/64), n_tok, 1]

---

追踪 4：GPU Flash Attention

Prefill 注意力

// 行 10569: Prefill 阶段的 Flash Attention
int ds4_gpu_attention_prefill_raw_heads_tensor(
        ds4_gpu_tensor *heads,         // 输出: [n_head × head_dim]
        const void *model_map,
        uint64_t model_size,
        uint64_t sinks_offset,         // Sink 参数偏移
        const ds4_gpu_tensor *q,       // Query [n_tok × n_head × head_dim]
        const ds4_gpu_tensor *raw_kv,  // KV 缓存 [n_tokens × head_dim]
        uint32_t n_tokens,             // KV 缓存行数
        uint32_t window,               // SWA 窗口大小
        uint32_t n_head, uint32_t head_dim) {

    // 包装 Sink 参数（零拷贝）
    ds4_gpu_wrap_model_range(model_map, sinks_offset, ...);

    // 获取专用 Flash Attention 管线
    ds4_gpu_encode_flash_attention_prefill_raw_heads(
        cb, pipe, heads, q, raw_kv,
        n_tokens, window, n_head, head_dim);
}

Decode 注意力

// 行 10619: Decode 阶段的 Flash Attention
int ds4_gpu_attention_decode_raw_batch_heads_tensor(
        ds4_gpu_tensor *heads,
        const void *model_map, uint64_t model_size,
        uint64_t sinks_offset,
        const ds4_gpu_tensor *q,
        const ds4_gpu_tensor *raw_kv,
        uint32_t n_tokens,
        uint32_t pos0,                 // 序列起始位置
        uint32_t n_raw,                // raw 缓存行数
        uint32_t raw_cap,              // raw 缓存容量
        uint32_t raw_start,            // raw 缓存起始索引
        uint32_t window,
        uint32_t n_head, uint32_t head_dim) {

    // 额外参数管理 KV 缓存窗口:
    // pos0:       当前 token 在完整序列中的位置
    // n_raw:      实际有效的 raw KV 行数
    // raw_cap:    raw 缓存的总容量
    // raw_start:  循环缓冲区的起始索引

    ds4_gpu_encode_flash_attention_decode_raw_batch_heads(
        cb, pipe, heads, q, raw_kv,
        n_tokens, pos0, n_raw, raw_cap, raw_start, ...);
}

Flash Attention GPU 实现:

标准注意力:               Flash Attention:
  1. Q×K^T → scores       1. 分块加载 Q, K, V 到共享内存
  2. softmax(scores)      2. 在线 softmax（分块计算最大值和指数和）
  3. × V → output         3. 分块累加输出
  内存: O(n²)             4. 不显式存储完整的 attention 矩阵
                          内存: O(n)

Prefill: 所有 prompt token 的 Q 对所有 KV 做注意力
Decode:  1 个新 token 的 Q 对所有 KV 做注意力

---

追踪 5：GPU MoE 与 HC 操作

路由 MoE

// 行 12736: GPU 上的路由 MoE
int ds4_gpu_routed_moe_one_tensor(
        ds4_gpu_tensor *out,
        ds4_gpu_tensor *gate, ds4_gpu_tensor *up,
        ds4_gpu_tensor *mid,  ds4_gpu_tensor *experts,
        const void *model_map, uint64_t model_size,
        uint64_t gate_offset, uint64_t up_offset,
        uint64_t down_offset,
        uint32_t gate_type, uint32_t down_type,
        uint64_t gate_expert_bytes, uint64_t gate_row_bytes,
        uint64_t down_expert_bytes, uint64_t down_row_bytes,
        uint32_t expert_in_dim, uint32_t expert_mid_dim,
        uint32_t out_dim,
        const ds4_gpu_tensor *selected,    // 6 个专家索引
        const ds4_gpu_tensor *weights,     // 6 个专家权重
        uint32_t n_expert,                 // 6
        float clamp,
        const ds4_gpu_tensor *x) {         // 输入

    // 对齐要求:
    // expert_in_dim  % 256 == 0
    // expert_mid_dim % 256 == 0

    // 验证 8 个 Metal 缓冲区大小
    // gate, up, mid, out, experts, selected, weights, x

    // GPU 内核执行:
    // 1. 6 个专家的 gate+up 投影（并行）
    // 2. SwiGLU 激活
    // 3. 6 个专家的 down 投影（并行）
    // 4. 加权累加
}

HC Split Sinkhorn

// 行 13503: GPU 上的 HC 分解
int ds4_gpu_hc_split_sinkhorn_tensor(
        ds4_gpu_tensor *out,
        const ds4_gpu_tensor *mix,         // 混合系数
        const void *model_map,
        uint64_t model_size,
        uint64_t scale_offset,             // 缩放参数偏移
        uint64_t base_offset,              // 基础偏置偏移
        uint32_t n_hc,                     // HC 维度
        uint32_t sinkhorn_iters,           // Sinkhorn 迭代次数
        float eps) {

    // 计算输出大小
    uint32_t mix_hc = 2 * n_hc + n_hc * n_hc;  // pre + post + comb
    // out 大小 = mix_hc × sizeof(float)

    // 验证 n_hc ∈ [1, 16]

    // 包装模型范围（零拷贝）
    ds4_gpu_wrap_model_range(model_map, scale_offset, ...);
    ds4_gpu_wrap_model_range(model_map, base_offset, ...);

    // 获取管线并分发
    id<MTLComputePipelineState> pipe =
        ds4_gpu_get_pipeline("hc_split_sinkhorn");
}

---

追踪 6：GPU 命令提交与同步

命令缓冲区管理

// 行 3903: 开始新的命令批次
int ds4_gpu_begin_commands(void) {
    g_batch_cb = [g_queue commandBuffer];       // 从队列获取命令缓冲区
    return 0;
}

// 行 3910: 刷新（提交当前批次，立即开始新批次）
int ds4_gpu_flush_commands(void) {
    [g_batch_cb commit];                        // 提交到 GPU
    [g_pending_cbs addObject:g_batch_cb];       // 追踪待完成命令
    g_batch_cb = [g_queue commandBuffer];        // 立即创建新批次
    return 0;
}

// 行 3929: 结束命令批次
int ds4_gpu_end_commands(void) {
    [g_batch_cb commit];
    [g_pending_cbs addObject:g_batch_cb];
    g_batch_cb = nil;
    return 0;
}

// 行 3937: 等待所有 GPU 操作完成
int ds4_gpu_synchronize(void) {
    if (g_batch_cb) ds4_gpu_end_commands();     // 先结束当前批次
    for (id<MTLCommandBuffer> cb in g_pending_cbs)
        [cb waitUntilCompleted];                // 等待每个命令完成
    [g_pending_cbs removeAllObjects];
    return 0;
}

GPU 命令流水线:

CPU 时间线:
  begin → 编码命令 → flush → 编码更多 → flush → end → synchronize
          │              │                │          │        │
          ▼              ▼                ▼          ▼        ▼
GPU 时间线:
       [执行批次1]    [执行批次2]     [执行批次3]  [等待完成]

flush 的作用:
  - 提交已编码的命令，让 GPU 立即开始执行
  - 同时创建新的命令缓冲区继续编码
  - CPU/GPU 重叠工作，隐藏延迟

synchronize:
  - 等待所有批次完成
  - 通常在需要读取 GPU 输出时调用
  - 例如: 采样前需要 logits 已经写好

---

追踪 7：CUDA Prefill 性能优化

宽 WMMA 评分 Kernel

CUDA 后端使用 NVIDIA WMMA（Warp Matrix Multiply-Accumulate）张量核心进行压缩索引评分：

// 原有: 16 压缩 token/tile (16x16 WMMA)
indexer_scores_wmma_kernel

// 新增变体（默认使用 128 宽）:
indexer_scores_wmma32_kernel   // 32/tile, 64 threads
indexer_scores_wmma64_kernel   // 64/tile, 128 threads
indexer_scores_wmma128_kernel  // 128/tile, 256 threads

// 优化: 压缩索引预加载到 shared memory，只加载一次
// 然后迭代所有注意力头，避免每头重复读全局内存

CUB BlockRadixSort Top-K

// 替代手写 bitonic sort，处理 n_comp <= 8192
// 512 threads × 16 items = 8192 items per block
indexer_topk_8192_cub_kernel

// Float → uint64_t key packing 实现降序排序:
// key = (topk_float_ordered_key(v) << 32) | (0xffffffff - idx)
// 整数降序 = float 降序

// 动态 shared memory:
cudaDevAttrMaxSharedMemoryPerBlockOptin  // 查询设备上限
cudaFuncSetAttribute(..., shared_size)   // 请求额外 shared memory

Top-K 预排序优化访存

// 将 512 个 top-k 索引按升序排列后再送入注意力 kernel
indexed_topk_sort_512_asc_kernel
// 随机访问 → 顺序访问，改善压缩 KV cache 的内存合并

实测：32K prefill 255→346 tok/s（+36%），全曲线均值 316→370 tok/s（+17%）。

---

追踪 8：Metal NAX 加速

Neural Accelerator 检测

// ds4_metal.m — M5 NAX 检测
// 启动时检查系统属性中的 M5 标志
g_metal4_m5_neural_accelerators_hint = ...;  // 读取设备属性

// NAX kernel 选择逻辑：
if (g_metal4_m5_neural_accelerators_hint && n_tokens >= 16) {
    // 使用 NAX 加速路径
    // indexer_scores: 卸载到 Neural Accelerator
    // MoE expert matmul: Q8_0 → NAX tensor matmul
} else {
    // 标准金属计算路径
}

NAX kernel 自动在 M5 设备上启用，无需用户配置。在非 M5 设备上编译安全（条件分支指向标准路径）。

---

总结：ds4 代码架构全景

ds4.c (18189 行) — 推理引擎核心
├── 量化内核 (158-3990)        量化与矩阵运算
│   ├── block_q2_K/q4_K/q8_K/iq2_xxs
│   ├── matvec_q8_0 / matvec_q2_k / matvec_iq2_xxs
│   └── f16_to_f32 / fp8 量化
├── 线程池 (612-770)           推理流程
├── 注意力机制 (2700-7092)     模型架构
│   ├── RMS Norm / RoPE        推理流程
│   ├── Q/KV 投影 / Flash Attention
│   ├── KV 缓存管理
│   └── 混合注意力 (raw + comp)
├── MoE / FFN (5012-6031)      模型架构
│   ├── SwiGLU 激活
│   ├── Hash/Top-K 路由
│   └── 共享专家 + 路由专家
├── 推理循环 (6055-7919)       推理流程
│   ├── Prefill (层主序)
│   └── Decode (逐 token)
├── 分词器 (13438-14165)       分词与采样
├── 采样 (14874-15083)         分词与采样
└── 公共 API (15716-17440)     模型架构

ds4_server.c (10349 行) — HTTP 服务
├── JSON 解析器 (141-437)      服务层
├── HTTP/SSE (3080-3180)       服务层
├── API 兼容层 (735-4259)      服务层
├── KV 持久化 (5138-6195)      服务层
├── SHA-1 (5192-5270)          服务层
└── 服务器架构 (4717-8100)     服务层

ds4_metal.m (14647 行) — GPU 后端
├── 初始化与管线 (624-2661)    GPU 加速
├── 张量管理 (3767-3937)       GPU 加速
├── 矩阵运算 (4919-5160)       GPU 加速
├── Flash Attention (10569+)   GPU 加速
├── MoE / HC (12736+)          GPU 加速
└── Metal Shader (.metal)       编译到 .metallib