CUDA 后端
colibrì 的默认路径是纯 CPU + 磁盘流式。它还有一个可选的 CUDA 层(-DCOLI_CUDA),把最热的固定专家和常驻稠密张量放进 GPU 显存。核心设计取舍是:流式专家故意留在 CPU 路径上——每次用到都从 NVMe 拷到 GPU 只会把磁盘瓶颈换成 PCIe 瓶颈。只有一次上传、反复复用的常驻张量才值得上 GPU。实现在 backend_cuda.cu 和 glm.c 的集成代码。
1. 设备侧库 backend_cuda.cu
backend_cuda.cu(230 行)是一个自包含的 extern "C" CUDA 库,它对 GLM 一无所知——只是持有量化 [O,I] 张量的常驻设备副本并跑一个量化矩阵乘 kernel。
ColiCudaTensor(8–14 行):不透明的设备副本(weights、per-row scales、fmt/I/O/device)。DeviceContext(16–21 行):每 GPU 的可复用输入/输出缓冲 + 张量计数。静态数组g_ctx[COLI_CUDA_MAX_DEVICES](23 行),最多 16 个设备。row_bytes(41 行):按 fmt 算每行字节(f32/int8/int4/int2)。__device__ weight_at(49 行):设备端反量化单个元素(int4 减 8、int2 减 2)——与 CPU 内核一致的语义。__global__ quant_matmul(62 行):kernel。grid =(O, S),256 线程/块,每块用 256 宽共享内存树归约(73 行)算一个输出行,量化时乘 scales[o](81 行)。
公开 C ABI(都在 backend_cuda.h):
| 函数 | 行 | 用途 |
|---|---|---|
coli_cuda_init | 94 | 校验设备序号,每 GPU 开一个 context,打印名称/显存/SM |
coli_cuda_mem_info | 142 | cudaMemGetInfo,用于 VRAM 预算 |
coli_cuda_tensor_upload | 158 | 一次上传权重(+scales),已常驻则只校验 |
coli_cuda_matmul | 191 | 惰性上传后跑 kernel,拷回结果 |
coli_cuda_tensor_free | 210 | 释放设备张量 |
它是完全同步的 scratch-stream 设计:每次调用一次阻塞 H2D 输入拷贝 → kernel → D2H 拷回。README 明说这些是”correctness-first”的自定义 kernel,不是 cuBLAS/Tensor Core,还没做端到端加速声明。
2. 主机侧集成
QT(69–75 行) 带了 CUDA 字段:ColiCudaTensor *cuda(72 行) 和 cuda_eligible/cuda_failed/cuda_device(74 行)——注释强调这些只给常驻张量,永不给复用的流式专家 slot。
全局状态(146–151 行,都在 #ifdef COLI_CUDA 内):g_cuda_enabled(147)、g_cuda_expert_gb(148)、g_cuda_dense(149)、g_cuda_devices[16](150)、g_cuda_dense_projected[16](151)。
助手:qt_cuda_upload(156 行)(按 fmt 选权重指针上传)、parse_cuda_devices(169 行)(解析 "0,1,2",拒绝重复/溢出)。
3. CUDA 层如何与磁盘流式分层组合
CUDA 是叠加在 CPU+磁盘基础之上的可选顶层,由三块可组合的部分构成:
3.1 分派点
matmul_qt(456–470 行):只有当 g_cuda_enabled && w->cuda_eligible && !w->cuda_failed && !omp_in_parallel()(462 行) 时才走 GPU。任何 CUDA 失败就置 cuda_failed 并回退 CPU(465–468 行)。!omp_in_parallel() 这个 guard 让嵌套的逐专家 OpenMP 工作留在 CPU——流式 LRU 专家 slot 永远不上 GPU。
3.2 可选稠密层
qt_load(690 行):g_cuda_dense 时把每个常驻稠密张量标 cuda_eligible=1,round-robin 分配设备,字节累加进 g_cuda_dense_projected。稠密张量在首次矩阵乘时惰性上传。
3.3 热专家(固定)层
固定例程内(2226–2278 行):专家在 RAM 固定后,若 g_cuda_expert_gb>0,算出每设备余量 = free − dense_projected − 2 GB,把请求预算钳到安全值,然后按热度顺序把最热的固定专家的 gate/up/down 上传到能装下的最空设备。上传失败就降级回 RAM(2264–2268 行)。所以一个固定专家是 RAM 常驻且额外镜像到 VRAM。
3.4 回合间 repin
repin_pass(1818 行)(见 07)换冷固定专家为热的(磁盘→RAM 进原 slot);若该 slot 是 GPU-backed,立即刷新同一 VRAM slot(1837–1849 行),失败降级 RAM。
4. 环境变量接线
| 环境变量 | 作用 |
|---|---|
COLI_CUDA=1 | 启用 CUDA 层 |
COLI_GPU / COLI_GPUS | 单卡 / 多卡设备列表 |
CUDA_DENSE=1 | 启用稠密张量上 GPU |
CUDA_EXPERT_GB | 固定专家的 VRAM 预算(多卡时是总预算) |
没有 COLI_CUDA=1 却设了这些旋钮会被 guard 拒绝。多卡时 CUDA_EXPERT_GB 是设备集的总预算,专家整体分配到能装下的最空设备。
5. 当前限制(README 明示)
- 设备用独立 context + 同步的 host-staged 激活拷贝,没有 P2P/NCCL。
- kernel 是 correctness-first 的自定义 kernel,不是 cuBLAS/Tensor Core。
- 流式专家故意不上 GPU(避免 PCIe 瓶颈换磁盘瓶颈)。
- NUMA-local 的 RAM backing store 还没实现。
README 提供了一个 313M 参数的确定性 glm_moe_dsa fixture(make_glm_bench_model.py),用于在没有完整 checkpoint 时做后端 A/B(CPU 流式 vs 稠密 CUDA vs CPU 热存储 vs CUDA 热专家)。
6. 小结
| 机制 | 出处 | 设计取舍 |
|---|---|---|
| 设备侧量化矩阵乘 | quant_matmul:62 | 自定义 kernel,correctness-first |
| 只给常驻张量 | matmul_qt:462 | 流式专家永不上 GPU |
| 可选稠密上 GPU | qt_load:690 | CUDA_DENSE=1,惰性上传 |
| 热专家镜像 VRAM | pin:2226 | RAM 常驻 + VRAM 镜像,失败降级 |
| 多卡总预算分配 | parse_cuda_devices:169 | 整体分配到最空设备 |