第 8 章 · 02 CUDA/Metal/Vulkan 三后端

第 8 章 · 02 CUDA/Metal/Vulkan 三后端

本节摘要:本节深潜 Colibrì 的三个可选 GPU 后端——backend_cuda.cu(NVIDIA CUDA,也可编 HIP)、backend_metal.mm(Apple Metal)、backend_vulkan.c(跨平台 Vulkan,主攻 Strix Halo 集成显存)。三者共享同一套函数指针契约(第 8 章 01 节),但取舍完全不同:CUDA 强调离散显卡专家驻留、Metal 强调统一内存零拷贝、Vulkan 强调"在 iGPU 上让流式专家 offload 划算"。本节还会看 c/shaders/ 下的四个 .comp 文件,以及"预编译 spv vs 运行时编译"的取舍。

内容来源:原项目源码 c/backend_cuda.cu、c/backend_metal.mm、c/backend_vulkan.c、c/shaders/*.comp

⚠️ 注意:三个后端各自有编译开关 CUDA=1 / METAL=1 / VULKAN=1,三者互斥(也可全不选纯 CPU)。每个后端都是一份独立源文件 + 一组 shader,代码量巨大(单 backend_metal.mm 就把整个 Metal shader 内嵌为字符串常量)。本节只贴头部注释和关键设计,不展开 kernel 内部。

学习目标

1. 读懂三个后端各自的"目标硬件 + 核心策略"。
2. 区分 CUDA 的"专家驻留"、Metal 的"统一内存零拷贝"、Vulkan 的"iGPU 流式专家 offload"。
3. 认识四个 GPU shader:attention_absorb.comp / qmatmul.comp / qmatmul_gate_up.comp / rmsnorm.comp。
4. 解释为什么 qmatmul.spv 预编译而其他 shader 运行时编译。
5. 理解 GPU 专家驻留如何减少 decode 期磁盘 IO。

一、backend_cuda.cu:CUDA/HIP 双 vendor 后端

backend_cuda.cu 顶部说明它服务的硬件和工作模式:

1   #include "backend_cuda.h"
2
3   #include "backend_gpu_compat.h"
...
17   struct RaggedKVEntry {
18       const void *key;
19       const float *host_l,*host_r;
20       float **latent_pages,**rope_pages;
21       int length,page_count,K,R;
22   };
24   #ifndef COLI_KV_PAGE_TOKENS
25   #define COLI_KV_PAGE_TOKENS 64
26   #endif
​

CUDA 后端的核心策略可归纳为三点(细节散在源码、docs/ 与 GPU_BACKENDS.md):

• 专家驻留(expert residency):把高频专家权重上传到显存常驻,decode 时直接 GPU 算,不再每 token 从磁盘拉。配置项 CUDA_EXPERT_GB=auto + PIN_GB 控制驻留预算——预算满则全驻留,磁盘读取退出主路径。
• page 化 KV:COLI_KV_PAGE_TOKENS=64 表明 KV latent/rope 缓存按 64 token 一页分配,decode 时按绝对位置追加。这和 kv_persist.h 的位置索引一致。
• HIP 复用:#if defined(__HIP_PLATFORM_AMD__) || defined(__HIP__) 一段表明这份 .cu 用 HIP 编译就变成 AMD 后端,沿用 coli_cuda_ 符号前缀。

此外还有 backend_cuda_ink.cu——多模态引擎 inkling.c 的 CUDA 变体,处理音频帧特殊路径。

💡 深潜要点:CUDA 后端的设计假设是"离散显卡 + PCIe":权重得 H2D 拷贝过去才能算,所以策略是"把热点专家一次性搬上去常驻",而不是"每 token 现搬"。这与下面 Vulkan 的 iGPU 路线截然相反。

二、backend_metal.mm:Apple 统一内存零拷贝

backend_metal.mm 顶部 4 行说尽核心:

1   // Apple-GPU (Metal) backend for colibrì. Runtime-compiled shader (no Xcode needed),
2   // zero-copy over unified memory. See backend_metal.h and docs/plans/2026-07-10-*.
​

• 运行时编译 shader:Metal shader 字符串内嵌在 .mm 里(就是源码里那个 static const char *SHADER = R"METAL(...)METAL";),运行时交给 Metal framework 编译。不需要 Xcode 工程——一行 clang 命令就能构建。
• 统一内存零拷贝:Apple Silicon 是 unified memory,CPU 和 GPU 共享同一物理 DRAM。权重不需要 H2D 拷贝,GPU 直接读 CPU 写的内存。
• shader 内部直接还原 quant.h 的精度表:E8G[256] 这种常量数组是从 quant.h 的 e8_grid 生成的,注释明确指出"metal-test oracle 把 256 个字节码全跑一遍 CPU 解码器比对,任何漂移会让构建失败"。这就是 token-exact 验证在 shader 层的落地。

实验演进路径(见 SUMMARY.md):早期 E4 用 MTLHeap 分配大块 slab,macOS 15.0+ 之后切到 MTLResidencySet(常驻集)做更精细的驻留管理。源码里有 macOS 15.0+ guard 保护新 API 调用,旧系统退化。

三、backend_vulkan.c:跨平台 + iGPU 流式专家 offload

backend_vulkan.c 的开场注释把它的"目标硬件"和"为什么不学 CUDA"讲得极清:

1   // Vulkan compute backend for colibri's quantized matmul, targeting the
2   // Strix Halo iGPU (RADV gfx1151). Mirrors backend_cuda.c's contract but
3   // exploits unified memory: weight "uploads" write into HOST_VISIBLE|
4   // DEVICE_LOCAL memory — the same physical RAM the iGPU reads — so there is
5   // no PCIe copy. That is what makes offloading *streamed experts* profitable
6   // here, which the discrete-CUDA path deliberately avoids.
...
9   // M2 scope: correctness + a standalone GPU-vs-CPU test harness. Synchronous
10  // submit/wait per call; async queues and zero-copy import come in M4.
​

要点:

• 目标硬件是 Strix Halo 集成 GPU(RADV gfx1151,AMD APU);
• 权重"upload"写到 HOST_VISIBLE|DEVICE_LOCAL 内存——这是同一块物理 RAM,GPU 直接读,没有 PCIe 拷贝;
• 正因为无拷贝,"流式专家 offload"在 iGPU 上才划算——而离散 CUDA 路径"刻意避免"流式专家(每 token 走 PCIe 搬运不划算);
• M2 阶段是"正确性 + GPU-vs-CPU 测试台",同步 submit/wait;异步队列和零拷贝 import 留给 M4。

ColiVkTensor 结构(注释下方)记录 wbuf/sbuf 两个 VkBuffer、对应 VkDeviceMemory、以及格式 fmt(0=f32/1=i8/2=i4/...)、I/O(输入/输出维度)、gs(group size)。还内置一个 VK_KV_LAYERS=160 的持久化设备 KV latent/rope 缓存——和 CUDA 的 kv_dev shadow 一样,host 按绝对位置追加行,absorb kernel 原地读。

四、c/shaders/:四个 GLSL compute shader

Vulkan 路径的 shader 是独立 .comp 文件,放在 c/shaders/:

attention_absorb.comp 头部注释把整个核心算法写在注释里:

1   #version 450
2   // MLA weight-absorption attention core (decode), the VK equivalent of
3   // coli_cuda_attention_absorb / glm.c's CPU absorb loop. ...
10   //   qabs[i]  = sum_d q_nope[d] * dequant(Wkv_b[h*(Q+V)+d])[i]        (absorbed query)
11   //   score_t  = (qabs . L_t + q_rope . R_t) * scale     t in [st0, T-S+s]  (causal)
12   //   p        = softmax(score)
13   //   clat[i]  = sum_t p_t * L_t[i]                                    (latent context)
14   //   ctx[v]   = clat . dequant(Wkv_b[h*(Q+V)+Q+v])                    (value rows)
​

四步:吸收 query(absorb)→ 算注意力分数 → softmax → 还原 value。L/R 是持久化的设备 KV latent/rope 缓存,host 每 token 追加一行。

qmatmul.comp 头部注释强调它"镜像 colibri 的 CUDA quant_matmul",并明说优化技巧参考自 llama.cpp 的 mul_mat_vec(MIT):

1   #version 450
2   // Quantized mat-vec/mat-mat mirroring colibri's CUDA quant_matmul:
3   //   y[s,o] = ( sum_i x[s,i] * dequant(W[o,i]) ) * scale[o]
...
9   // Optimized decode (S=1) GEMV, techniques adapted from llama.cpp's mul_mat_vec
10   // (MIT) to colibri's per-row int4/int8 layout:
11   //   - x[s,:] loaded ONCE into shared memory, reused across every output row ...
​

五、预编译 spv vs 运行时编译

c/shaders/ 下既有 .comp(源码)也有 qmatmul.spv(SPIR-V 二进制)。区别:

• qmatmul.spv 预编译:Vulkan 后端启动时直接加载 SPIR-V 二进制,跳过编译期,首次启动快、避免依赖 shader 编译器(glslangValidator 等);
• 其他 shader 运行时编译:如果环境里有编译器就源码编译,适合调试期改 shader 后立即看效果。

这个取舍对应 Colibrì 一贯的"热路径稳定 + 实验路径灵活":被验证过、不变的热 shader 预编译进仓;还在迭代的 shader 留源码。

六、GPU 专家驻留:decode 期省磁盘 IO

把第 8 章前两节串起来——GPU 专家驻留的目标不是"算得更快",而是"decode 期不读磁盘":

• 流式 MoE 引擎 decode 时,每 token 的路由专家若在 NVMe 上,要现读——这是磁盘 IO 主导的成本(第 5 章、第 7 章 02 节实测);
• 如果把高频专家提前 upload 到显存(或 Apple unified memory、或 Strix Halo 的 HOST_VISIBLE DEVICE_LOCAL 内存)并常驻,decode 命中时直接 GPU 算,磁盘读退出主路径;
• 这正是 CUDA_EXPERT_GB=auto + PIN_GB、MTLResidencySet、Vulkan iGPU offload 三条路线的共同目标。

注意 Colibrì 反复强调的边界:这只是"减少 decode 期磁盘 IO",不是"GPU 一定比 CPU 快"。是否真省时,取决于 PCIe 拷贝 / 统一内存读取 / 算力 / 驻留预算的综合权衡(第 8 章 03 节展开)。

本节要点回顾

• backend_cuda.cu 服务离散 NVIDIA(可编 HIP),核心策略 = 专家驻留(CUDA_EXPERT_GB=auto + PIN_GB)+ KV page 化 + inkling 变体 backend_cuda_ink.cu。
• backend_metal.mm 服务 Apple Silicon,统一内存零拷贝 + shader 字符串内嵌运行时编译(免 Xcode);从 MTLHeap(E4)演进到 MTLResidencySet(E5,macOS 15.0+)。
• backend_vulkan.c 主攻 Strix Halo iGPU,利用 HOST_VISIBLE|DEVICE_LOCAL 让"流式专家 offload"在 iGPU 上划算——离散 CUDA 刻意避免这条路。
• 四个 shader:attention_absorb.comp(MLA 注意力吸收)/qmatmul.comp(量化矩阵乘)/qmatmul_gate_up.comp(门控上投影)/rmsnorm.comp(RMS 归一化)。
• qmatmul.spv 预编译(热路径稳定),其他 shader 运行时编译(实验路径灵活);GPU 专家驻留的目标是 decode 期省磁盘 IO,不是 GPU 必快。

下一节把第 8 章收口:CPU/GPU 重叠隐藏的是传输而非移动瓶颈,OpenMP 调优只动线程数不动 spin-wait,以及"快 CPU + 低驻留可能抹平 GPU 收益"的诚实假设。