Candle项目中的FP16矩阵乘法优化探讨

Candle是一个基于Rust的深度学习框架，近期社区成员发现其在处理FP16精度的矩阵乘法运算时存在性能优化空间。本文将深入分析这一技术问题，并探讨解决方案。

问题背景

在深度学习推理过程中，矩阵乘法(GEMM)是最核心的计算操作之一。当使用FP16精度时，NVIDIA CUDA提供了两种计算方式：

1. 使用FP16进行累加(hgemm)
2. 使用FP32进行累加(sgemm)

社区成员通过性能分析工具发现，Candle当前使用的是FP32累加方式(turing_fp16_s1688gemm_fp16_256x128_ldg8_f2f_tn内核)，而同类框架如llama.cpp则使用了FP16累加方式(turing_h1688gemm_256x128_ldg8_tn内核)。实测显示，后者能带来约15%的性能提升。

技术分析

精度与性能的权衡

FP16累加的优势在于：

• 更高的计算吞吐量
• 更低的内存带宽需求
• 更快的执行速度

但同时也存在：

• 累加过程中可能丢失精度
• 数值稳定性稍差

FP32累加则相反，虽然速度稍慢，但能保持更高的计算精度。

Candle的实现细节

Candle通过cudarc库调用CUDA BLAS的通用GEMM接口，默认使用CUBLAS_COMPUTE_32F计算类型。这意味着虽然输入输出都是FP16，但累加过程使用FP32精度，这与PyTorch的默认行为一致。

解决方案

经过讨论和测试，Candle项目采取了以下改进措施：

1. 保留了FP32累加作为默认选项，确保大多数模型的数值稳定性
2. 新增了全局标志，允许用户显式启用FP16累加模式
3. 特别针对量化模型场景优化了默认配置

这种设计既照顾了通用场景的稳定性需求，又为追求极致性能的特殊场景提供了选择权。

实际效果

在Mistral.rs项目中的测试表明，启用FP16累加后：

• 推理速度提升了约15%
• 从1000 token/s提升到1150 token/s
• 在量化模型上未观察到明显的精度损失

行业实践参考

PyTorch也提供了类似的精度控制选项：

• 默认允许降低精度以获得性能
• 但可通过torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction进行配置

这一设计思路与Candle的解决方案不谋而合，体现了深度学习框架在精度与性能间寻求平衡的通用做法。

总结

Candle项目通过这次优化，展示了其灵活性和对性能的追求。对于开发者而言，理解框架底层计算细节有助于更好地调优应用性能。特别是在量化模型等对精度要求不苛刻的场景，适当降低计算精度可以带来可观的性能提升。

未来随着硬件和算法的发展，如何在各种场景下智能选择最优计算精度，仍是深度学习框架需要持续探索的方向。