Como é que o Flash Attention reduz o tráfego de memória durante a inferência local?

O Flash Attention reduz o tráfego de memória ao dividir a atenção exata em blocos, permitindo processar as pontuações intermédias no próprio chip em vez de as escrever e ler repetidamente na HBM da GPU.

Num servidor de IA local, o efeito torna-se mais importante à medida que o comprimento do prompt aumenta, porque a atenção padrão cria uma quantidade cada vez maior de trabalho intermédio. O Flash Attention não altera o conhecimento do modelo nem elimina a respetiva cache KV; reorganiza o cálculo da atenção para mover menos dados através dos níveis mais lentos da hierarquia de memória da GPU.

A Atenção Padrão Move Grandes Tensores Intermédios Através da HBM

A atenção começa com tensores de consulta, chave e valor, e tem de calcular as pontuações, normalizá-las e combiná-las com os valores. Uma implementação direta materializa grandes matrizes intermédias na memória de elevada largura de banda entre essas etapas.

O artigo original sobre FlashAttention defende que o desempenho da atenção é limitado não apenas pela aritmética, mas também pelas leituras e escritas entre a HBM e a SRAM no chip. As sequências longas amplificam esse tráfego, porque a matriz de pontuações cresce com os pares de tokens.

Uma GPU doméstica pode, por isso, passar uma quantidade significativa de tempo a mover dados em vez de executar operações úteis com matrizes, especialmente quando o comprimento do contexto aumenta.

O Flash Attention Divide o Problema em Conjuntos de Trabalho no Chip

O Flash Attention processa blocos de consultas, chaves e valores dimensionados para caber na memória rápida no chip. Cada bloco executa várias etapas da atenção antes de os dados intermédios saírem do chip.

O FlashAttention-3 descreve a abordagem central da família como uma estratégia de divisão em blocos que evita o tráfego intermédio na memória global. O algoritmo continua a produzir atenção exata para a formulação suportada.

O tamanho dos blocos depende dos recursos do hardware e da implementação do kernel. O mecanismo importante é a localidade dos dados: reutilizar os valores carregados enquanto estes permanecem próximos das unidades de cálculo.

O Softmax Online Evita Armazenar a Matriz de Pontuações Completa

O softmax normal parece exigir todas as pontuações de uma linha antes da normalização, o que obrigaria a implementação a materializar uma matriz de grandes dimensões. Em vez disso, o Flash Attention mantém estatísticas de normalização acumuladas ao longo dos blocos.

O tutorial de atenção fundida do Triton demonstra atenção fundida com estado de softmax por blocos. Isto permite consumir e descartar blocos parciais de pontuações sem escrever todos os valores intermédios novamente na HBM.

É aqui que a redução do tráfego de memória e a redução da memória temporária se encontram: o algoritmo reorganiza o cálculo para que um objeto intermédio de grandes dimensões nunca tenha de existir na memória global.

Featured

ZimaCube 2 NAS Pessoal Cloud para Casa

ZimaCube2

$799.00 - $2,499.00

Buy Now

A Fusão de Kernels Mantém uma Maior Parte da Cadeia de Atenção num Único Percurso de Execução

O cálculo das pontuações, o escalamento, a aplicação da máscara, a normalização e a agregação dos valores podem ser fundidos num kernel coordenado, em vez de serem lançados kernels separados que trocam todos os dados intermédios através da memória global.

A interface de atenção com produto escalar escalado do PyTorch pode selecionar implementações de atenção fundidas quando o dispositivo, o tipo de dados, a forma e as funcionalidades solicitadas são suportados.

O resultado pode ser um processamento mais rápido dos prompts e uma menor utilização máxima de memória temporária, mas o ganho exato depende de o ambiente de execução realmente encaminhar o trabalho para um kernel compatível com Flash Attention.

A Redução das Operações de E/S Não É o Mesmo que Eliminar a Cache KV

Durante a inferência autoregressiva, as chaves e os valores gerados anteriormente continuam a ter de permanecer disponíveis para os tokens seguintes, a menos que o modelo ou o ambiente de execução utilize um mecanismo separado de remoção, compressão, janela ou transferência.

A documentação da ROCm da AMD descreve o Flash Attention como uma forma de melhorar a localidade da memória entre o armazenamento no chip e a HBM. Essa otimização aborda a forma como a atenção lê e calcula o estado ativo, não a existência do estado KV histórico.

Esta distinção evita um erro comum de planeamento: o Flash Attention pode reduzir a memória de trabalho e as operações de E/S, enquanto um servidor com contexto longo continua a ficar sem capacidade porque a cache KV retida continua a crescer.

O Hardware e a Forma dos Dados Determinam se o Percurso Rápido é Utilizado

Os kernels Flash são especializados. A dimensão da cabeça, o tipo de dados, a máscara de atenção, a arquitetura da GPU, a versão do ambiente de execução e outras funcionalidades podem obrigar a utilizar uma implementação alternativa.

O artigo da ZimaSpace sobre memória da atenção face aos parâmetros do modelo aborda a questão da capacidade. O Flash Attention deve ser medido separadamente como uma otimização do percurso de execução.

Faça testes com prompts curtos e longos, confirme o backend selecionado e registe a latência do prefill, a memória alocada máxima, a largura de banda da memória e a taxa de tokens de saída. O guia de hardware de IA local da ZimaSpace fornece a base mais ampla de capacidade que a otimização de kernels não pode substituir.

O Transformer Engine da NVIDIA aborda backends de atenção otimizados e execução dependente do hardware como parte de uma pilha mais ampla de seleção de kernels.