新PyTorch API：幾行代碼實現不同注意力變體，兼具FlashAttention性能和PyTorch靈活性

AIGC動態歡迎閱讀

原標題：新PyTorch API：幾行代碼實現不同注意力變體，兼具FlashAttention性能和PyTorch靈活性
關鍵字：注意力,變體,內核,性能,因果
文章來源：機器之心
內容字數：0字

內容摘要：

機器之心報道
編輯：陳陳用 FlexAttention 嘗試一種新的注意力模式。理論上，注意力機制就是你所需要的一切。然而在實際操作中，我們還需要優化像 FlashAttention 這樣的注意力機制的實現。
盡管這些融合的注意力機制大大提高了性能，且支持長上下文，但這種效率的提升也伴隨著靈活性的喪失。對于機器學習研究人員來說，這就像是一種「軟件彩票」—— 如果你的注意力變體不適合現有的優化內核，你將面臨運行緩慢和 CUDA 內存不足的困境。
一些注意力變體包括因果注意力、相對位置嵌入、Alibi、滑動窗口注意力、PrefixLM、文檔掩碼、不規則張量、PagedAttention 等。更糟糕的是，人們通常希望將這些變體組合在一起！比如滑動窗口注意力 + 文檔掩碼 + 因果注意力 + 上下文并行，又比如 PagedAttention + 滑動窗口的組合。
下圖左側代表了當今的現狀 —— 一些掩碼 + 偏置 + 設置的組合已經有現成的內核實現。然而，各種選項的添加會導致設置呈指數級增長。更糟糕的是，這種方式不會支持新的注意力變體。為了徹底地解決這個超立方體問題，PyTorch 團隊引入了

原文鏈接：新PyTorch API：幾行代碼實現不同注意力變體，兼具FlashAttention性能和PyTorch靈活性

聯系作者

文章來源：機器之心
作者微信：
作者簡介：