characterizing Mamba-based 状态空间模型在 GPU 上的训练行为
机器学习
2025-08-26 v1 硬件体系结构
计算与语言
摘要
Mamba-based 状态空间模型（SSM）已成为颠覆了普遍使用的 transformer 的有前景的替代方案。尽管 transformer 拥有强大的表达能力，但注意力计算的二次复杂度是扩大序列长度时性能瓶颈所在。SSM 提供了一条不同路径，通过针对不同领域和应用（如视频、文本生成和图形）的新型模型架构，降低了自注意力的计算复杂度要求。因此，重要的是在 GPU 上对这些新兴工作负载进行特征化，并理解其在 GPU 微架构设计中的需求。在本工作中，我们对 Mamba-based SSM 在 GPU 上进行训练评估并进行特征化。我们构建了一个覆盖不同模型架构的工作负载套件。随后，我们利用该套件分析在 GPU 上运行 Mamba-based SSM 的架构影响。我们的工作为继续扩展此类模型的性能提供了新的视角，揭示了潜在的优化方向。
引用
@article{arxiv.2508.17679,
title = {Characterizing the Behavior of Training Mamba-based State Space Models on GPUs},
author = {Trinayan Baruah and Kaustubh Shivdikar and Sara Prescott and David Kaeli},
journal= {arXiv preprint arXiv:2508.17679},
year = {2025}
}