Finalmente, depois de dias lendo documentos ou depurando índices de thread, consegui implementar a atenção instantânea do zero no DSC no MI300X! A primeira versão (laranja) é a versão escalar básica do artigo original de flash-attention. A segunda versão (verde) é o mesmo algoritmo, mas usa núcleos de matriz (núcleos tensores AMD) e, como você pode ver, isso é * significativamente * mais rápido que o escalar. Usei os núcleos da matriz para calcular Sij = Qi @ Kj ^ T e Pij @ Vj. Algumas 'pegadinhas' dos núcleos de matriz AMD: - Eles funcionam por frente de onda e uma frente de onda é de 64 threads na AMD, isso significa que você precisa acompanhar o ID da onda atual e também o ID do thread dentro dessa onda. - O layout de saída será embaralhado nos registros devido ao fato de que a operação principal de um núcleo de matriz é um produto externo 4x1, portanto, é necessária uma etapa de reordenação. - (Até onde eu sei) os intrínsecos do hipcc para núcleos de matriz não estão documentados em nenhum lugar. Há um repositório com vários exemplos da AMD, mas fora isso você terá que grep a base de código LLVM. Vou polir meu código de vez em quando e provavelmente escreverei um post mais aprofundado sobre a atenção do flash na AMD. Ah, e aliás, grite para @HotAisle por tornar isso possível!