CUDA原生最小LLM推理：剥离框架后的内存布局与核函数调度

当 PyTorch 和 Transformers 库封装了几乎所有底层细节，开发者往往对 Transformer 在 GPU 上的真实执行路径缺乏直观理解。纯 CUDA 实现的最小语言模型推理框架提供了一个剥离抽象层、直面硬件特性的实验场。这类项目以零依赖、单文件或极简代码结构为目标，让开发者能够直接操控内存布局、量化格式和核函数调度策略。

为什么需要手写 CUDA 实现

现有推理框架如 vLLM、TensorRT-LLM 虽然性能优异，但代码复杂度高、依赖繁重，难以用于教学目的和快速原型验证。相比之下，calm 等项目的核心目标是 "最大单 GPU 单 batch 硬件利用率" 与 "最小实现" 的平衡。这类实现通常只有数千行 CUDA 代码，却支持 Llama、Mistral、Qwen2、Mixtral 等主流 decoder-only 架构。

手写 CUDA 的核心价值在于理解自回归推理的带宽瓶颈本质：每生成一个 token 都需要读取整个模型权重和完整的 KV 缓存。当使用 fp16 精度时，7B 参数模型仅权重读取就需要约 14GB 内存带宽；加上 KV 缓存访问，推理速度几乎完全受限于 GPU 内存带宽而非算力。

内存布局与量化策略

权重存储格式选择

最小 CUDA 实现通常支持三种量化格式，在模型大小、推理速度和精度之间权衡：

• fp16（e5m10）：基准精度，无精度损失但带宽占用最高
• fp8（e5m2）：8 位浮点，约 0.5% 的困惑度损失，推理速度接近翻倍
• gf4（分组 4 位）：8 个值用 32 位存储（3 位量化 scale 加 fp8 组 scale），速度提升约 75%，模型大小减半

实际测试显示，在 RTX 4090 上运行 Llama3 8B 时，fp16 约 61 tok/s，fp8 可达 120 tok/s，gf4 则达到 225 tok/s。值得注意的是，gf4 量化是均匀且纯粹的 —— 所有层都精确量化为 4 位每权重，这与 llama.cpp 的 K-quants 策略不同。

KV 缓存的动态精度切换

KV 缓存默认使用 fp16 存储，但在长上下文场景（>4096 tokens）下，实现会自动切换至 fp8 以节省内存并提升性能。这种动态切换机制需要在核函数中处理精度转换开销，但带来的内存节省通常能抵消转换成本。

核函数调度与单 batch 优化

矩阵乘法的内存访问模式

自回归推理的核心是矩阵乘法：Q/K/V 投影、注意力计算、FFN 前馈。在纯 CUDA 实现中，开发者需要手动设计线程块和网格布局，确保全局内存合并访问（coalesced access）。

关键优化点包括：

• 权重布局：采用行优先或列优先存储以匹配 CUDA 的 warp 访问模式
• 共享内存复用：在注意力计算中复用 Q/K/V 的共享内存缓冲区
• 寄存器压力控制：避免过度展开循环导致寄存器溢出到本地内存

带宽利用率最大化

在 RTX 4090 上，calm 实现了约 900-950 GB/s 的有效带宽，接近理论峰值 1008 GB/s。达到这一利用率的关键在于：

1. 消除同步点：最小化核函数间的__syncthreads()调用
2. 预取指令：使用 LDG 指令预取下一轮迭代所需数据
3. 计算与内存重叠：在支持异步拷贝的架构上利用 CUDA 流

对于 Mixtral 8x7B MoE 模型，在 H100 SXM 上使用 fp8 权重可达到约 200 tok/s，有效带宽约 2550 GB/s。

可 Hack 的代码结构

最小实现的教育价值在于其可修改性。典型的代码结构包括：

src/
├── model.cu      # 模型加载与前向传播主循环
├── kernels.cu    # CUDA核函数实现（矩阵乘、注意力、激活函数）
├── quantize.cu   # 量化/反量化核函数
└── main.cu       # 命令行接口与采样逻辑

开发者可以便捷地修改以下组件：

• 采样策略：在 main.cu 中调整 temperature、top-p、top-k 参数
• 注意力变体：在 kernels.cu 中替换标准 attention 为 FlashAttention 或线性注意力
• 量化方案：在 quantize.cu 中实验 e4m3 fp8 或 INT8 量化

实际可落地的参数配置

基于实测数据，以下是不同场景下的推荐配置：

对于上下文长度超过 4096 的场景，建议显式启用 KV 缓存 fp8 模式，并监控困惑度退化情况。

局限与扩展方向

当前最小 CUDA 实现的主要局限包括：

• 仅支持单 batch 推理：多 batch 并行需要额外的内存管理和核函数调度
• Prompt 并行处理缺失：提示词仍需逐 token 串行处理，未利用并行前缀计算
• 量化校准简化：缺乏逐层敏感度分析和自适应量化策略

扩展方向可考虑：

1. 集成 FlashAttention-2 以降低长序列的内存复杂度
2. 实现投机采样（speculative decoding）以突破内存带宽瓶颈
3. 添加多 GPU 张量并行支持

资料来源

• zeux/calm: CUDA/Metal accelerated language model inference, GitHub
• karpathy/llama2.c: Inference Llama 2 in one file of pure C, GitHub

ai-systems