最后，在几天的阅读文档或调试线程索引后，我终于能够在MI300X上从头实现DSC中的闪存注意力！ 第一个版本（橙色）是来自原始闪存注意力论文的基本标量版本。 第二个版本（绿色）是相同的算法，但使用矩阵核心（AMD张量核心），正如你所看到的，这比标量版本*显著*快。 我使用矩阵核心计算Sij = Qi @ Kj^T和Pij @ Vj。 一些AMD矩阵核心的“注意事项”： - 它们在每个波前的基础上工作，而一个波前在AMD上是64个线程，这意味着你必须跟踪当前波的ID以及该波内线程的ID。 - 由于矩阵核心的核心操作是4x1外积，输出布局将在寄存器中被打乱，因此需要一个重新排序步骤。 - （据我所知）矩阵核心的hipcc内在函数没有任何文档。AMD有一个包含许多示例的仓库，但除此之外，你必须grep LLVM代码库。 我现在要润色我的代码，然后可能会写一篇关于AMD闪存注意力的更深入的文章。 哦，顺便说一下，感谢@HotAisle让这一切成为可能！