最後，在幾天的閱讀文檔或調試線程索引後，我終於能夠在MI300X上從頭實現DSC中的閃存注意力！ 第一個版本（橙色）是來自原始閃存注意力論文的基本標量版本。 第二個版本（綠色）是相同的算法，但使用矩陣核心（AMD張量核心），正如你所看到的，這比標量版本*顯著*快。 我使用矩陣核心計算Sij = Qi @ Kj^T和Pij @ Vj。 一些AMD矩陣核心的“注意事項”： - 它們在每個波前的基礎上工作，而一個波前在AMD上是64個線程，這意味著你必須跟蹤當前波的ID以及該波內線程的ID。 - 由於矩陣核心的核心操作是4x1外積，輸出佈局將在寄存器中被打亂，因此需要一個重新排序步驟。 - （據我所知）矩陣核心的hipcc內在函數沒有任何文檔。AMD有一個包含許多示例的倉庫，但除此之外，你必須grep LLVM代碼庫。 我現在要潤色我的代碼，然後可能會寫一篇關於AMD閃存注意力的更深入的文章。 哦，順便說一下，感謝@HotAisle讓這一切成為可能！