marlin - 专为LLM推理设计的FP16xINT4优化内核

Marlin 项目介绍

项目概览

Marlin 是一款名为混合自回归线性核（Mixed Auto-Regressive Linear kernel）的高效计算核心，同时也以地球上最快的鱼类之一命名。这款核心专注于大语言模型（LLM）推理，其优化的 FP16xINT4 矩阵乘法操作能够在推理时实现接近理想的 4 倍加速，适用于 16-32 个 token 的批次处理，这相比于之前相似加速效果却仅支持 1-2 个 token 的情况，有了显著的提升。Marlin 的这种特性使其非常适合于大规模服务、推测性解码或是复杂的多重推理方案。

关键技术

Marlin 运用了多种技术来同时优化利用 GPU 的全球内存、L2 缓存、共享内存、张量核心和矢量核心，确保性能提升能够在实践中实现：

优化缓存使用：激活函数的计算方式保证多次从寄存器而非共享内存中读取，减轻共享内存的负担。此外，异步执行全局权重加载并立即清除未使用值，避免污染 L2 缓存。
内存加载与计算并行化：采用双缓冲技术将共享内存的加载与全局加载、计算重叠进行，提高运行效率。
高效的指令排序：仔细排列去量化和张量核指令以确保 GPU 管道的饱和度，避免瓶颈。
结构化的数据布局：在执行时以理想的访问模式组织量化权重和组缩放值，促进数据的高效处理。
计算结果聚合与内存优化：通过 "条纹" 分区方案提高计算单元利用率，并在输出缓冲区中直接进行全局化简，结合静态内存偏移，最大化计算效率并最小化冗余读取和写入。

性能表现

Marlin 在 NVIDIA A10 GPU 上测试，与其他流行的 4-bit 推理核心相比，展示了极其强大的性能优势。Marlin 在所有批大小下几乎都达到理想的 3.87 倍加速。即使在较小的矩阵和不同型号的 GPU 上，Marlin 通过其独特的分区策略，仍然保持了突出的速度表现。

长时间运行的性能研究表明，Marlin 即使在基础时钟速度下，其优越性能也未受影响，而其他核心在这种情况下则性能减弱明显。

使用要求

CUDA 版本 11.8 或更高（尤其对于 nvcc 编译器，同步 torch 版本）
计算能力 >= 8.0 的 NVIDIA GPU（Ampere 或 Ada 平台）
torch>=2.0.0
numpy
完整量化流程需要：transformers、datasets 和 sentencepiece

安装与使用

满足上述要求后，可通过以下命令安装 Marlin：

pip install .

安装后，可通过 marlin.Layer（一个表示 Marlin 量化层的 torch 模块）使用 Marlin 核心。此外也可直接通过 marlin.mul(..) 调用核心，前提是权重和缩放值已经适当预处理。相关代码位于 marlin/marlin_cuda_kernel.cu 文件中，不依赖于基础 CUDA 之外的环境，方便集成到其他底层框架中。

核对与测试

可通过以下命令进行正确性测试和基准性能测试：

python test.py  # 进行正确性测试
python bench.py # 进行基准测试

留意：为重现可持续性能的基准测试，需要使用以下命令将 GPU 时钟锁定至基础值：

sudo nvidia-smi --lock-gpu-clocks=BASE_GPU_CLOCK --lock-memory-clocks=BASE_MEM_CLOCK

此外，在 A10 GPU 等开启 ECC 的情况下，最大可实现的内存带宽将比官方规格降低10%-15%，可通过以下命令禁用 ECC：

sudo nvidia-smi -e 0

GPTQ 示例

在 gptq 子目录下，提供了改进的 GPTQ 算法版本，改善了组栅格截断和非均匀校正样本长度，支持生成兼容 Marlin 的 4-bit Llama2 模型。此外，还有脚本可以用来在流行的 LLM 评估工具中评估这种压缩模型。