Tri Dao团队开发出完全基于Python的内存密集型GPU内核优化库QuACK,使用CuTe-DSL实现,无需CUDA代码即可在H100上达到近“光速”吞吐量。实测显示,其softmax内核吞吐量达3.01 TB/s,较torch.compile快近50%。该方案通过硬件感知的加载/存储策略与多层级归约机制,有效利用Hopper架构的线程块集群特性,在大规模归约任务中显著优于现有优化库。英伟达与PyTorch专家对此给予肯定,并探讨性能调优方向。作者指出未来或可通过LLM自动生成高效GPU内核。
来源:https://mp.weixin.qq.com/s/ZL_COJl1m-V83iDlC6BmDw