Appearance
Part 6: GPU 加速
Metal、CUDA、Flash Attention
1. GPU 后端 + 总结
CPU decode 速度约 5-10 t/s,对话体验卡顿。GPU 后端将矩阵运算卸载到 Metal/CUDA/ROCm,decode 速度提升 5-10 倍。今天学习 ObjC 互操作、Metal 计算管线和 Flash Attention,最后用一个完整的请求生命周期串联 15 天的知识。
GPU 后端将 CPU 注意力 卸载到 Metal/CUDA/ROCm,大幅提升推理速度。
C 知识点
1. Objective-C 互操作
ds4_metal.m 用 Objective-C 调用 Metal API,ds4_cuda.cu 用 CUDA C++ 调用 NVIDIA API。ds4.c 是纯 C,通过 ds4_gpu.h 统一抽象层与 GPU 后端交互:
c
// ds4_gpu.h — GPU 后端无关的抽象接口
// 所有函数签名使用 ds4_gpu_tensor * 不透明指针
// 不暴露 Metal (MTLBuffer) 或 CUDA (CUdeviceptr) 类型
typedef struct ds4_gpu_tensor ds4_gpu_tensor;
// 张量生命周期
ds4_gpu_tensor *ds4_gpu_tensor_alloc(ds4_gpu_graph *g, const char *name, uint64_t nbytes);
void ds4_gpu_tensor_free(ds4_gpu_tensor *t);
void ds4_gpu_tensor_read(ds4_gpu_tensor *t, void *dst, uint64_t offset, uint64_t len);
void ds4_gpu_tensor_write(ds4_gpu_tensor *t, const void *src, uint64_t offset, uint64_t len);
// 命令队列
void ds4_gpu_begin(ds4_gpu_graph *g);
void ds4_gpu_flush(ds4_gpu_graph *g);
void ds4_gpu_end(ds4_gpu_graph *g);
void ds4_gpu_synchronize(ds4_gpu_graph *g);
// ~60 个推理内核: embedding, RoPE, attention, MoE, steering 等
void ds4_gpu_embedding(ds4_gpu_graph *g, ...);
void ds4_gpu_flash_attn(ds4_gpu_graph *g, ...);
// ...后端枚举 (ds4.h):
c
typedef enum {
DS4_BACKEND_METAL,
DS4_BACKEND_CUDA,
DS4_BACKEND_CPU,
} ds4_backend;编译选择:
- macOS: 默认 Metal 后端 (
CORE_OBJS = ds4.o ds4_metal.o) - Linux: 默认 CUDA 后端 (
CORE_OBJS = ds4.o ds4_cuda.o,用nvcc链接) make cpu: 两平台通用,-DDS4_NO_GPU编译纯 CPU 版本- CUDA 构建目标:
make cuda-spark: DGX Spark / GB10 专用(无需指定 arch)make cuda-generic: 通用 CUDA GPUmake cuda CUDA_ARCH=sm_N: 显式指定架构(如sm_80)
- ROCm 构建目标(AMD Strix Halo /
gfx1151):make strix-halo(别名make rocm):用hipcc链接ds4_rocm.o+rocm/*.cuh,定义-DDS4_ROCM_BUILD,默认ROCM_ARCH=gfx1151- ROCm 后端现已并入 main 分支(此前长期只在独立的
rocm分支,由社区因 antirez 无对应硬件而维护)
- 质量评分工具
gguf-tools quality-score也支持跨平台:macOS 用$(CC)链接 Metal,Linux 用$(NVCC)链接-lcudart -lcublas,通过#ifdef __APPLE__编译时选择后端
关键设计:
ds4_gpu.h对 C 完全隐藏具体 GPU API 类型(Metal/CUDA/HIP)- 新增 GPU 后端只需实现
ds4_gpu.h中的 ~60 个函数,不改动ds4.c - CPU 路径通过
DS4_NO_GPU宏完全绕过 GPU 代码 - ROCm 后端靠编译期
-DDS4_ROCM_BUILD选择(而非运行期 enum 值):它在ds4_gpu.h里用#ifdef DS4_ROCM_BUILD暴露若干独有钩子(layer-major 批量加载、q8_f16缓存释放等),其它后端提供 no-op stub 保持链接。ds4_backend枚举仍是METAL/CUDA/CPU——ROCm 复用了非 Apple(Linux)的代码路径,只是把实现文件从ds4_cuda.cu换成ds4_rocm.cu。
2. Metal 计算管线
Metal GPU 执行模型推理的核心流程:
ds4_metal.m 中的流程:
objc
id<MTLDevice> device = MTLCreateSystemDefaultDevice();
// 启动时打印设备名称和内存
// ds4: Metal device Apple M2 Ultra, 192.00 GiB RAM
id<MTLCommandQueue> queue = [device newCommandQueue];
// 编译 .metal shader
id<MTLLibrary> lib = [device newLibraryWithSource:shader_source ...];
id<MTLFunction> fn = [lib newFunctionWithName:@"kernel_name"];
id<MTLComputePipelineState> pipeline = [device newComputePipelineStateWithFunction:fn ...];
// 执行
id<MTLCommandBuffer> cmd = [queue commandBuffer];
id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
[enc setComputePipelineState:pipeline];
[enc setBuffer:input offset:0 atIndex:0];
[enc setBuffer:output offset:0 atIndex:1];
[enc dispatchThreadgroups:... threadsPerThreadgroup:...];
[enc endEncoding];
[cmd commit];
[cmd waitUntilCompleted];可选 buffer 绑定占位
encoder 绑定 buffer 时,可选参数即使不使用也必须绑定占位值,否则 Metal debug layer 会报校验错误:
objc
// ds4_metal.m — router finalize
float zero_f32 = 0;
if (has_bias)
[enc setBuffer:bias_buf offset:0 atIndex:3];
else
[enc setBytes:&zero_f32 length:4 atIndex:3]; // 零值占位
if (hash_mode)
[enc setBuffer:hash_buf offset:0 atIndex:5];
else
[enc setBytes:&zero_f32 length:4 atIndex:5]; // 零值占位Metal debug layer 要求 shader 声明的每个 [[buffer(N)]] 都有对应绑定,不用的可选参数用 setBytes:&zero 填充即可。
3. Metal Shading Language
.metal 文件用 MSL 编写 GPU kernel:
ds4 项目有 19 个 .metal 文件,覆盖:
dense: 矩阵乘法flash_attn: Flash Attentionmoe: MoE 专家计算norm: RMSNormsoftmax: Softmaxdsv4_kv/dsv4_rope: KV cache 和 RoPEargsort: 排序(用于 top-k)- 等等
Metal Neural Acceleration (NAX)
M5 芯片引入了 MPP/Neural Accelerator 硬件(NAX),ds4 在检测到 M5 时自动利用两个 NAX kernel:
- Indexer Scores NAX — 压缩 KV 注意力的索引评分,将评分计算卸载到 Neural Accelerator
- Tensor Matmul NAX — MoE 专家的 Q8_0 矩阵乘法,支持 32/64/128 tile sizes
检测方式:g_metal4_m5_neural_accelerators_hint(从系统属性读取的 M5 标志)。当 n_tokens >= 16 时启用 NAX 路径,低于此阈值时退回标准 GPU 计算。
Compressed KV in F16
Metal 后端新增 comp_kv_f16 模式——将压缩 KV cache 以半精度(float16)而非 float32 存储。每行 KV 内存减半(从 DS4_N_HEAD_DIM × 4B 降为 DS4_N_HEAD_DIM × 2B),在大上下文场景下节省大量内存。对应的 Metal kernel 位于 metal/dsv4_kv.metal,在读入时将 F16 解压缩为 F32 用于注意力计算。
GPU 功耗节流
GPU 后端支持 power_percent(1-100)参数控制推理速度。实现方式是在每个 prefill 层和 decode token 后调用 graph_power_sleep():
sleep_us = (100 - power) / power × elapsed_usagent 模式下通过 ds4_engine_set_power() 动态调整——等待用户确认时降低功耗,活跃推理时恢复满速。
Typed Pointer vs void Pointer
Metal shader 参数应使用 device const char * 而非 device const void *,让 Metal 反映正确的只读访问语义:
get_rows.metal 和 set_rows.metal 等文件中的参数已从 void * 改为 char *。void * 在 Metal 中不传达内存访问意图,char * 配合 const 明确表达只读语义。
LLM 知识点
1. GPU 加速原理
CPU vs GPU 的根本区别:
CPU: 4-12 个强核心,每个核心处理复杂任务
GPU: 数千个弱核心,每个核心做简单计算但并行执行
矩阵乘法 A[n,m] × B[m,k]:
CPU: 逐元素计算,串行(或少量并行)
GPU: 每个元素一个线程,数千元素同时计算为什么 LLM 推理需要 GPU?
- 矩阵乘法是高度并行的(每个输出元素独立)
- GPU 的内存带宽远高于 CPU
- Metal GPU 可以直接引用 mmap 内存(零拷贝)
2. Flash Attention
标准注意力:先算完整的 Q·K^T 矩阵,再 softmax。内存 O(n²)。
Flash Attention:分块计算,每次只加载一小块 Q 和 K 到高速缓存,在片上完成 softmax。内存 O(n)。
ds4 的 metal/flash_attn.metal 实现了这个优化。
3. 零拷贝 MTLBuffer
c
// Metal 可以直接引用 mmap 的内存,不需要拷贝到 GPU 专用内存
id<MTLBuffer> buf = [device newBufferWithBytesNoCopy:
(void*)(mmap_base + tensor_offset)
length:tensor_bytes
options:MTLResourceStorageModeShared
deallocator:nil];这就是为什么模型加载用 MAP_SHARED——Metal GPU 直接从文件映射读取权重,不需要额外的 GPU 内存拷贝。81GB 的模型只需要 81GB 的磁盘空间和 mmap 映射,不需要额外的 GPU 显存。
模型视图与张量覆盖不变量
当模型文件超过 Metal 设备的 maxBufferLength 时,引擎创建多个重叠的 MTLBuffer 视图来覆盖整个文件。相邻视图必须有足够大的重叠区域,确保每个张量至少完整包含在一个视图中:
c
// 最大单张量大小限制,决定视图重叠区域
#define DS4_METAL_MODEL_MAX_TENSOR_BYTES (2ull * 1024ull * 1024ull * 1024ull) // 2 GiB
// 重叠 = round_up(MAX_TENSOR_BYTES, page) + page
// 视图步进 = max_buffer - overlapQ4_K 量化模型的 MoE 专家张量约 1.125 GiB,之前的 672 MiB 限制不足以覆盖。2 GiB 阈值确保所有量化格式(包括 Q4 专家张量)都能完整映射到至少一个视图中。
4. 分布式推理(Distributed Inference)
PRO 模型(~430GB IQ2_XXS 或 ~838GB Q4 分片)太大,单机无法运行。ds4 新增 coordinator/worker 分布式推理架构,将模型层切片分配到多台机器上执行:
核心设计:
| 组件 | 职责 |
|---|---|
| Coordinator | 接收请求、执行前半层、转发激活、接收 logits |
| Worker | 注册层范围、执行后半层、返回结果 |
| 传输协议 | 自定义 TCP 二进制协议(magic 0x44533444 = "DS4D") |
| KV 快照 | 拓扑无关:保存时聚合所有 worker 的层张量,加载时按当前路由分发 |
Worker 通过 HELLO 消息注册(包含 layer_start、layer_end、model_id、quant_bits),Coordinator 据此构建路由表。激活传输支持可配置的 activation_bits(默认 32-bit float,可压缩为更低精度减少网络开销)。
bash
# Coordinator(加载前半层)
./ds4 --role coordinator -m pro-q4-layers00-30.gguf -p "Hello"
# Worker(加载后半层)
./ds4 --role worker -m pro-q4-layers31-output.gguf模型分片加载(Model Span Loading):分布式引擎使用 ds4_gpu_set_model_map_spans() API,只加载分配给本机的层范围,而非整个模型。这通过 offset/size 对描述需要映射的字节范围,配合 chunked preloading 异步拷贝到设备内存。
c
// 只加载本 worker 负责的层切片
ds4_gpu_set_model_map_spans(
model_map, model_size,
offsets, sizes, span_count,
max_tensor_bytes
);详见 分布式推理。
4a. SSD Streaming(SSD 流式推理)
除了分布式推理,ds4 还提供另一种突破内存限制的方式:SSD Streaming。核心洞察是 MoE 层有 256 个路由专家,但每个 token 只激活 6 个——路由专家占模型大部分空间,但大部分时间都在"睡觉"。
缓存计划:ds4_ssd_cache_plan 根据可用 RAM 预算自动计算能缓存多少专家。自动预算取推荐工作集的 80%、扣除非路由权重,剩下的用于路由专家;显式 --ssd-streaming-cache-experts 32GB 则直接给字节预算(换算成完整专家数),过大时推理前自动封顶以保持可锁定(lockable)。详见 SSD Streaming 术语表。
热度排名(启动):ds4_streaming_hotlist.inc(13,334 行)按流行度预排列所有 (layer, expert_id) 对。启动时预加载最热门的专家,自动预热默认封顶 4096 个。
运行期淘汰(路由热度):决定"淘汰谁"的是一张运行期热度表 route_hotness[layer][expert],每处理 16 个 decode token 整表右移衰减一位(>>= 1),淘汰时驱逐热度最低者。这张表是提示局部的——新 prompt 开始时清零(ds4_gpu_stream_expert_cache_reset_route_hotness()),但常驻缓存本身跨会话保持热:易变的启发式状态要重置,昂贵的物化缓存要复用。
三路 Prefill:默认 prefill 调度器按 prompt 长度分流——极短走逐 token decode 式、中等走 selected-expert 批量、长 prompt 走全层 layer-major。切分点按量化格式调优,避免短 prompt 落到"加载一堆用不上的专家"的慢路径。
混合精度路由专家:当 GGUF 跨层不统一量化(per-layer boosted quants,例如全局 IQ2_XXS 但有几层升到 Q4_K)时,提升层的 ffn_exps 张量被纳入 decode-span 构建器走 mapped-view 兜底,且 slab 分配器在启动时预置尺寸类、对超大尺寸 freeze+reject 而非 last-writer-wins,避免 slab 被毒化。
多后端实现:SSD streaming 不再是 Metal 独享——CUDA(ds4_cuda.cu,有界缓存 + 进度上报)和 ROCm(ds4_rocm.cu,selected-expert 缓存 + 重叠读取 + 全层 streaming prefill)都已实现。streaming API 已重构为传一张 ds4_gpu_stream_expert_table 表:
c
typedef struct ds4_gpu_stream_expert_table {
const void *model_map; uint64_t model_size;
uint32_t layer, n_total_expert;
uint64_t gate_offset, up_offset, down_offset;
uint64_t gate_expert_bytes, down_expert_bytes;
} ds4_gpu_stream_expert_table;
// 异步加载缺失专家(重构后参数更简洁)
int ds4_gpu_stream_expert_cache_begin_selected_load(
const ds4_gpu_stream_expert_table *table,
const int32_t *selected_ids, uint32_t n_selected);
uint32_t ds4_gpu_stream_expert_cache_current_count(void);命令行选项:
| 选项 | 说明 |
|---|---|
--ssd-streaming | 启用 SSD streaming 模式 |
--ssd-streaming-cache-experts N | 缓存字节预算(过大时推理前封顶) |
--ssd-streaming-preload-experts N | 显式指定启动预加载专家数(默认自动封顶 4096) |
--ssd-streaming-cold | 冷启动(不预加载 hotlist;仅用于测量) |
与分布式推理的区别:SSD streaming 是单机方案,适合内存不足但有大容量 SSD 的场景;分布式推理是多机方案,适合 PRO 模型等超出单机存储极限的情况。两者可以结合使用。
4b. DGX Spark / GB10 后端(HBM 缓存)
NVIDIA DGX Spark(GB10 芯片,121 GiB UMA)有一个特殊的 CUDA 推理挑战:ATS(Address Translation Service)允许 GPU 直接消费 host mmap 权重,但有效带宽远低于 HBM 驻留数据。
ds4 的解决方案:启动时 HBM 缓存——将"热"张量(注意力投影、MoE 共享专家、输出投影)拷贝到 GPU HBM,冷 MoE 路由专家保持 ATS 映射:
默认缓存预算 96 GiB(cuda_model_cache_limit_bytes()),通过 DS4_CUDA_WEIGHT_CACHE_LIMIT_GB 环境变量可覆盖。IQ2 模型(~81GB)和混合 Q2/Q4 模型(~91GB)可完整缓存;完整 Q4 模型(~153GB)超出预算,需使用分布式层加载。
HBM 缓存查找优先于 UVA 映射指针——直接 HBM 读取比通过 host 页表映射快约 10%。cuda_model_range_ptr() 使用 hash-keyed 查找实现单次读取命中。
效果(DGX Spark / GB10,ds4flash,n=256):
- 优化前:~13.9 t/s(纯 ATS 映射)
- 优化后:~16.13 t/s(HBM 缓存热张量 + 小 batch kernel 融合)
编译目标:make cuda-spark(自动检测 GB10 架构,无需指定 CUDA_ARCH)。
Metal View Cap 修复
分布式推理的 span 映射使用 ds4_gpu_add_model_view_range() 创建 Metal 视图。之前的实现给所有大映射都应用 128 GiB 视图上限,包括普通的单次完整模型映射,导致 Metal VM 验证变慢。
修复:新增 use_default_view_cap 参数——普通完整模型映射(ds4_gpu_map_model_views())不设上限,分布式 span 映射(ds4_gpu_set_model_map_spans())保留 128 GiB 分割。这样普通映射保持单次映射的高效性,只有需要切分的分布式场景才做视图分割。
模型缓存 FD 绑定
CUDA 后端使用直接文件 I/O(非 mmap)加载路由专家权重,需要正确的文件描述符和 mmap 基地址来计算文件偏移。当 MTP 模型在主模型 fd 绑定和缓存准备之间加载时,fd 可能被错误覆盖。
修复:新增 ds4_gpu_set_model_fd_for_map(int fd, const void *model_map) 接受显式的 fd 和 model_map 指针,替换隐式使用 g_model_host_base 的旧 ds4_gpu_set_model_fd()。MTP 模型设置完成后,重新绑定主模型的 fd:
c
// 缓存前绑定正确的 fd
ds4_gpu_set_model_fd_for_map(e->model.fd, e->model.map);
accelerator_cache_model_tensors(...);
// MTP 设置后重新绑定主模型 fd
ds4_gpu_set_model_fd_for_map(e->model.fd, e->model.map);4c. Strix Halo / ROCm 后端
AMD Strix Halo(如 Framework Desktop,Radeon 8060S,gfx1151,128 GB 统一内存)是 ds4 的第三个 GPU 后端。ROCm 后端现已并入 main 分支(此前长期只在独立的 rocm 分支,由社区维护,因为 antirez 没有对应硬件)。
构建:make strix-halo(别名 make rocm),用 hipcc 编译 ds4_rocm.o + rocm/*.cuh(约 18K 行新代码),定义 -DDS4_ROCM_BUILD。后端基于 rocWMMA(AMD 的矩阵乘加速库),需要 libhipblas/libhipblaslt/librocblas/librocwmma 等。
关键工程难点(详见上游 STRIXHALO.md):
ROCm 安装补全:Ubuntu 26.04 的
librocwmma-dev缺rocwmma/internal/头文件,需手动 clone rocWMMA 7.1.0 补齐;若 ROCm 装在/usr而工具链期望/opt/rocm,要建符号链接。显存扩展(GT TTM aperture):128 GB Strix Halo 系统默认只暴露约 62 GB GPU 可见内存,不够 80.76 GiB 模型 + 运行时缓冲。需设内核参数扩大 GTT:
amd_iommu=off amdgpu.gttsize=126976 ttm.pages_limit=32505856 ttm.page_pool_size=32505856重启后
rocminfo应报gfx1151 pool: 130023424 KB(约 124 GB)。GGUF 选择:在 Strix Halo 上用标准 IQ2XXS/Q2K/Q8 imatrix GGUF;避免 mixed IQ2/IQ4 或 IQ2/Q4 GGUF——它们给 ROCm 路径更大的内存压力,可能触发系统 OOM 而非干净的 ds4 失败。
ROCm 特定修复(并入本次同步):
Fix ROCm distributed slice mapping— 分布式层切片映射Fix ROCm MTP model residency— MTP 模型驻留Fix ROCm mixed streaming expert fallback— 混合精度 streaming 专家兜底
可选后端钩子(Fused GPU Ops as Backend Hooks):融合算子(如 fused RMSNorm/RoPE、fused MoE)从 ds4.c 里写死的 Metal/CUDA 分支重构为 ds4_gpu.h 上的可选后端钩子。每个后端按需实现(Metal 131 行、CUDA 94 行新增),ds4.c 侧减为统一的钩子分发(-40 行)。这让 ROCm 这类新后端能选择性提供融合内核而不必复制所有分支逻辑。
4d. 帮助系统(Help System)
新增结构化帮助模块 ds4_help.c/h,替代各工具中散乱的 usage 打印。API:
c
// ds4_help.h
typedef enum {
DS4_HELP_DS4, // ds4 CLI
DS4_HELP_SERVER, // ds4-server
DS4_HELP_AGENT, // ds4-agent
DS4_HELP_BENCH, // ds4-bench
DS4_HELP_EVAL, // ds4-eval
} ds4_help_tool;
void ds4_help_print(FILE *fp, ds4_help_tool tool, const char *topic);关键设计:
- 主题过滤:
topic参数可查询特定子主题(如--help distributed) - TTY 感知:
isatty()自动检测,终端中显示 ANSI 256-color 彩色输出,非终端输出纯文本 - 格式化助手:
opt()打印彩色选项行、para()打印黄色段落、title()打印章节标题 - 分布式模式指导:帮助文本中包含 distributed mode 的用法说明和示例
4e. 模型下载脚本更新
download_model.sh 新增 PRO 模型下载目标:
| 目标 | 说明 | 大小 |
|---|---|---|
pro-q2-imatrix | PRO IQ2_XXS 单文件 | ~430 GB |
pro-q4-layers00-30 | PRO Q4 前半层(coordinator) | ~426 GB |
pro-q4-layers31-output | PRO Q4 后半层(worker) | ~412 GB |
pro-q4-split | 下载两个 Q4 分片 | ~838 GB |
PRO 文件过大,脚本自动检测并使用 Hugging Face CLI(hf from huggingface_hub)下载,支持断点续传。未安装时提示 python3 -m pip install -U huggingface_hub hf_xet。
非 PRO 目标下载后自动符号链接到 ./ds4flash.gguf,PRO 目标打印分布式用法指引。
4f. 方向性激活引导(Directional Steering)
基于论文 Refusal in Language Models Is Mediated by a Single Direction 的思想,在推理时对每层激活做低秩编辑:
y = y - scale * direction[layer] * dot(direction[layer], y)- 引导文件:43 × 4096 的 f32 矩阵(43 层,每层一个 4096 维归一化方向向量)
- 可在 attention 输出后(
--dir-steering-attn)和/或 FFN 输出后(--dir-steering-ffn)施加 - 正 scale 移除该方向,负 scale 放大该方向
- 提取工具:
dir-steering/tools/build_direction.py(用好/坏两组 prompt 取激活差异,归一化后生成.f32文件) - 示例:
dir-steering/包含预构建的 verbosity 方向向量,负 FFN scale 使回答更简洁
4g. CUDA Q8→F16 缓存显存预算守卫
PRO 模型 GPU 注意:DeepSeek V4 PRO(1.6T 总参数,49B 激活参数)的专家张量更大(更多专家 × 更高维度),Q4 量化模型需要 512GB 内存的 Mac Studio。GPU 后端在加载 PRO 模型时会自动调整 KV cache 容量和 buffer 分配。
Q8 权重在推理时被扩展为 F16 存入 GPU 显存,加速后续 cuBLAS 矩阵乘法。但显存有限的 GPU 上,缓存可能占满 VRAM 导致后续分配失败:
c
// 显存预算检查流程
static int cuda_q8_f16_cache_has_budget(uint64_t request_bytes, const char *label) {
// 1. 检查用户设定的硬上限
const uint64_t limit = cuda_parse_mib_env("DS4_CUDA_Q8_F16_CACHE_MB");
// 2. 查询当前显存空闲量
cudaMemGetInfo(&free_bytes, &total_bytes);
// 3. 计算保留量:max(4 GiB, 5% VRAM)
reserve = max(4ull * 1024 * 1024 * 1024, total_bytes / 20);
// 4. 分配后剩余空间必须 >= reserve
return (free_bytes - request_bytes >= reserve);
}三层回退机制:
关键设计:缓存是纯加速路径,失败不影响正确性。环境变量控制:
DS4_CUDA_Q8_F16_CACHE_MB=0:完全禁用DS4_CUDA_Q8_F16_CACHE_MB=8192:限制缓存 8 GiBDS4_CUDA_Q8_F16_CACHE_RESERVE_MB=8192:覆盖默认保留量
4h. CUDA 长上下文修复
长上下文(>100K token)暴露了多个 CUDA kernel 的固定缓冲区假设:
Shared Memory 分数溢出 → 在线注意力 kernel 回退:
c
// 编译时常量限制了 shared memory 中的分数缓冲区大小
enum { DS4_CUDA_ATTENTION_SCORE_CAP = 8192u };
// 当压缩行数超过缓冲区容量时,切换到在线注意力 kernel
// 在线 kernel 不需要预分配所有分数,逐块处理多级 Tree-Merge TopK → 替代单层 chunk+merge:
c
// 旧:单层 chunk(2048) → merge,无法处理 >4096 压缩行
// 新:chunk(4096) → 多级 tree-merge,每 8 组合并一次
// n_sets: n_chunks → ceil(n/8) → ceil(n/64) → ... → ≤8 → 最终合并回归测试 tests/cuda_long_context_smoke.c:模拟长上下文场景(多个 prefill frontier + decode),验证 CUDA 路径在边界条件下的正确性。
MoE Down 路径清理
移除了 block16 MoE down 诊断路径(-172 行),因存在 top-logit 不稳定性和贪心首 token 行为差异。默认保留更快的 tile16/row2048 路径,block16 改为 DS4_CUDA_MOE_DOWN_BLOCK16 环境变量 opt-in。
CUDA 内存优化
三项改进减少不必要的 GPU 开销:compressor state 清零改用 cudaMemsetAsync(替代 fill kernel 调用);graph/session tensor 改用 cudaMalloc 设备内存替代 managed memory(避免多余的页面迁移);新增 backend fill_f32 hook 让 CUDA 设备端直接初始化张量。
CUDA Managed KV Cache
百万 token 级上下文时,KV cache 张量改用 cudaMallocManaged(统一内存按需分页),避免 DGX Spark 统一内存系统中 GPU 显存分配饿死 CPU 侧。普通上下文仍用 cudaMalloc 设备分配保持性能。通过 ds4_gpu_should_use_managed_kv_cache() 自动判断是否启用。
Metal Q4 Model View 回退
Metal Q4 expert tensor model views 的扩展(DS4_METAL_MODEL_MAX_TENSOR_BYTES 从 2 GiB 回到 ~672 MiB)已被回退——该修复由 AI 生成且未经用户确认,上游坚持不接受未经人工验证的修复。
4i. CUDA Prefill 性能优化
大幅优化长上下文 prefill 速度(32K prefill 255→346 tok/s,全曲线均值 316→370 tok/s),核心改动:
更宽的 WMMA 评分 kernel: 原有 indexer_scores_wmma_kernel 每 tile 处理 16 个压缩 token。新增 wmma32/64/128 变体(默认 128 宽,256 线程),将压缩索引预加载到 shared memory 后迭代所有注意力头,减少 kernel 启动开销并提高内存吞吐。
CUB BlockRadixSort 替代手写 bitonic sort: 新 indexer_topk_8192_cub_kernel 用 NVIDIA CUB 库的 BlockRadixSort(512 线程 × 16 items = 8192/block),将 float key 打包为 uint64_t 实现降序排序。通过 cudaFuncSetAttribute 请求额外 shared memory。
uint16_t 索引的 bitonic sort: 对 ≤65536 压缩 token 的场景,用 uint16_t 代替 uint32_t 索引,减半 shared memory 占用。
Top-K 预排序提升注意力访存: 新 indexed_topk_sort_512_asc_kernel 将 top-k 选出的 512 个索引按升序排列后再送入注意力 kernel,把对压缩 KV cache 的随机访问转变为顺序访问,显著改善内存合并。
Templatized 注意力 kernel: attention_indexed_mixed_heads8_online_kernel 重构为模板,参数化 ROWS_PER_STAGE 和 HEADS_PER_GROUP,启动配置从 <8, 8> 改为 <8, 16>(16 head/group, 512 线程),配合动态 shared memory。
Shared memory 复用: 原有 wmma kernel 交换加载顺序,压缩索引 tile 只加载一次(在头循环之前),避免每头重复读取全局内存。
c
// ds4.h — 引导选项
typedef struct {
// ...
const char *directional_steering_file;
float directional_steering_attn;
float directional_steering_ffn;
} ds4_engine_options;5. 增量吞吐量基准测试(ds4-bench)
传统基准测试报告全流程平均速度,ds4-bench 在不同 context frontier 处测量瞬时吞吐量:
- 每个 frontier(如 2048, 4096, 6144...)只计算新增 token 区间的 prefill 速度
- prefill 后保存内存快照,做固定 128 token 贪心解码测 generation 速度
- 快照保存/恢复时间不计入测量窗口
- 输出 CSV:每个 frontier 的 prefill t/s、generation t/s、kvcache_bytes
c
// ds4.h — 快照 API
typedef void (*ds4_session_progress_fn)(void *ud, int progress);
void ds4_session_set_progress(ds4_session *s, ds4_session_progress_fn fn, void *ud);
void *ds4_session_save_snapshot(ds4_session *s);
int ds4_session_load_snapshot(ds4_session *s, void *snapshot);
void ds4_session_snapshot_free(void *snapshot);测试数据使用 speed-bench/promessi_sposi.txt(23329 行意大利语公共领域文本)作为固定输入序列。bench/ 目录已重命名为 speed-bench/,新增 M2 Ultra (192 GB) 和 M4 Max 基准数据,以及 plot_speed.py CSV→SVG 可视化脚本。
15 天学习总结
C 语言技能树
LLM 推理知识树
ds4.c 的设计哲学
- 垂直整合:一个文件做所有事,零外部依赖
- 专用化:只服务一个模型,硬编码所有参数
- 零拷贝:mmap + Metal/CUDA no-copy buffer,数据永远不离开映射区
- 热路径零分配:预分配 scratch buffer,生成循环中不 malloc
- 早失败:加载时严格验证,不匹配就退出
- 可扩展后端:通过
ds4_gpu.h抽象接口,Metal/CUDA/ROCm/CPU 四后端可互换;融合算子作为可选后端钩子,新后端按需提供 - KV cache 是一等公民:渲染文本前缀匹配 + 对齐前沿 + KTM 工具记忆,让长上下文推理跨重启存活
- 方向性引导:推理时低秩编辑激活空间,无需重训练即可控制输出风格
- 分布式推理:coordinator/worker 模型切片执行,PRO 模型可跨多台机器运行
- PRO 模型支持:384 专家 Q4_K 量化路径,CPU/Metal/CUDA 三后端完整覆盖
- SSD Streaming:路由专家缓存 + SSD 按需加载 + 路由热度淘汰(每 16 token 衰减、提示局部重置但缓存跨会话保温)+ 三路 prefill,让模型在有限 RAM 下也能运行
- 合作式中断:Ctrl+C 在安全检查点优雅停止,KV cache 状态不丢失