dropbox/gemlite

解决的问题

GemLite 解决了 LLM 推理中低比特矩阵乘法的性能瓶颈。它提供了一组高度优化的内核，相比默认的 Torch AO 内核，显著加速了模型执行的预填充和解码阶段，降低了量化模型的延迟并提高了吞吐量。

工作原理

GemLite 使用 Triton 实现了多种针对不同工作负载优化的专用矩阵乘法内核：

• GEMM：基于标准 Tensor Core 的通用矩阵乘法内核。
• GEMM Split-K：通过将 K 维度拆分为多个任务来计算部分和，优化批量大小为 2-32 的批量 LLM 解码。
• GEMV：针对小批量（M=1）设计，将激活值拆分为一维块。
• GEMV RevSplit-K：一种新算法，通过将每个程序的工作量翻倍来优化批量大小为 1 的解码，从而减少加载缩放因子和零值的开销。

支持广泛的精度，包括 FP16、BF16、FP8、INT8 以及低比特权重（8、4、2、1 位），还支持 Blackwell 等新硬件的 MXFP 和 NVFP 格式。

适用人群

GemLite 面向致力于 LLM 推理优化的开发者和研究人员，特别是那些使用量化技术在保持 NVIDIA GPU 高性能的同时减小模型大小和内存带宽需求的人。

主要亮点

• 显著加速：相比默认的 Torch AO 内核，预填充速度提升高达 7-8 倍，解码速度提升高达 3-6 倍。
• 广泛精度支持：兼容 FP16、BF16、FP8、INT8 以及 MXFP/NVFP4。
• 集成性：可与 vLLM（通过 HQQ 或 TorchAO）、SGLang 集成，并支持 torch.compile。
• 自动调优缓存：支持保存和加载最优内核配置，消除启动时重复自动调优的时间开销。
• 灵活打包：支持 HQQ 风格的权重打包以及权重和激活的通道级缩放。