GEMM o Coração Matemático

A curiosa leitora precisa saber que a multiplicação de matrizes está no coração dos Transformers, mas a expressão multiplicar matrizes esconde quase toda a dificuldade prática. A definição matemática da multiplicação de matrizes cabe em uma linha. Isso não é suficiente. Qualquer implementação capaz de aproveitar uma CPU, ou GPU moderna, exige compreender hierarquias de memória, ordem dos acessos, blocagem, vetorização, paralelismo, precisão numérica e o custo de transportar cada byte. Cada byte conta.

Em álgebra linear de alto desempenho, uma multiplicação economizada pode ajudar. Um acesso evitado à memória pode ajudar ainda mais.

Nosso objeto de estudo será a operação GEMM, de General Matrix-Matrix Multiplication. A atenta leitora verá como a mesma equação atravessa diferentes camadas da máquina: memória principal, caches, registradores vetoriais, memória compartilhada da GPU e unidades matriciais especializadas. Ao final, poderemos olhar para um Transformer e reconhecer quais operações são limitadas por computação e quais são limitadas pelo movimento de dados. Isso, vai fazer toda a diferença nos seus projetos futuros. Uma vantagem competitiva em tempos de agentes de inteligência artificial.

1. Da multiplicação de matrizes à GEMM

Sejam $A$, $B$ e $C$ três matrizes tal que:

Se $C$ for a matriz produto, a multiplicação de matrizes clássica, das matrizes $A$ e $B$, define cada elemento da matriz $C$ como:

Vale a pena desmontar essa fórmula antes de generalizá-la.

O elemento $C_{ij}$ é o produto escalar entre a linha $i$ de $A$ e a coluna $j$ de $B$: percorremos as duas em paralelo, multiplicamos os pares de valores e somamos os $K$ produtos. Daí a exigência sobre as dimensões internas: a linha de $A$ tem $K$ elementos e a coluna de $B$ também, e sem essa coincidência o emparelhamento não existe. As dimensões externas, $M$ e $N$, apenas contam quantas linhas e quantas colunas há para combinar, e é delas que $C$ herda o formato $M\times N$.

Um exemplo pequeno deixa o mecanismo visível. Com $M=2$, $K=3$ e $N=2$:

O elemento $C_{00}$ combina a linha $0$ de $A$ com a coluna $0$ de $B$:

Repetindo o procedimento para as quatro posições:

Duas propriedades dessa definição sustentarão tudo o que veremos adiante, e ambas se provam em uma linha, olhando para um elemento genérico. Primeiro, multiplicar por um escalar comuta com o produto:

Segundo, a soma de matrizes de mesmo formato é definida elemento a elemento, $[X+Y]{ij}=X{ij}+Y_{ij}$; portanto, somar ao produto $AB$ qualquer outra matriz $M\times N$ é uma operação bem definida e barata, que não envolve nenhuma multiplicação de matrizes adicional. Guarde essas duas observações: elas são a licença algébrica para tudo o que a GEMM acrescenta.

Falta apresentar o personagem que usaremos o artigo inteiro: a transposta. A transposta de uma matriz $X\in\mathbb{R}^{m\times n}$, denotada $X^\top$, é a matriz em $\mathbb{R}^{n\times m}$ obtida trocando o papel de linhas e colunas:

No nosso exemplo, $B\in\mathbb{R}^{3\times2}$ tem transposta

Nenhum valor muda; muda apenas qual índice conta linhas e qual conta colunas. A transposta aparecerá com frequência nos Transformers, pois os scores de atenção da Seção 13 são $QK^\top$, e aparece nas bibliotecas por um motivo adicional, mundano: dependendo de como uma matriz está disposta na memória, é mais barato lê-la como se estivesse transposta do que fabricar fisicamente uma cópia transposta.

Com esses ingredientes, podemos enunciar a operação que as bibliotecas BLAS (Basic Linear Algebra Subprograms, Subprogramas Básicos de Álgebra Linear) de fato implementam:

na qual $\operatorname{op}(X)$ pode ser a matriz $X$ ou sua transposta $X^\top$, e $\alpha$ e $\beta$ são escalares. Essa operação é chamada GEMM, General Matrix Multiply (Multiplicação Geral de Matrizes). Alguns prefixos podem ser usados para indicar o tipo numérico: SGEMM para float, DGEMM para double e assim por diante. Elemento a elemento, chamando de $C^{\text{velho}}_{ij}$ o conteúdo de $C$ antes da operação, a GEMM calcula:

A atenta leitora tem o direito de exigir, neste ponto, a garantia de que nada de novo foi inventado. E não foi. O somatório no centro da fórmula é exatamente a definição de produto que abriu esta seção, aplicada às matrizes $\operatorname{op}(A)\in\mathbb{R}^{M\times K}$ e $\operatorname{op}(B)\in\mathbb{R}^{K\times N}$, cujo produto tem o formato $M\times N$ de $C$, como deve. Multiplicar esse produto por $\alpha$ e somar $\beta\,C^{\text{velho}}$ são as duas operações elemento a elemento que acabamos de legitimar. Em particular, escolhendo $\alpha=1$, $\beta=0$ e $\operatorname{op}$ como a identidade, a fórmula colapsa para

que é, símbolo por símbolo, a multiplicação clássica. A GEMM não é uma multiplicação diferente: é a multiplicação de sempre, embrulhada em uma atualização afim que a máquina consegue executar quase de graça na saída do produto.

Um exemplo numérico fecha o argumento. Reutilizando $A$ e $B$ acima, com $\alpha=2$, $\beta=1$ e um conteúdo prévio $C^{\text{velho}}$ igual à identidade:

Conferindo o primeiro elemento pela fórmula elemento a elemento: $2\,(1\cdot4+2\cdot2+0\cdot1)+1\cdot1=17$. O mesmo resultado seria obtido com três operações separadas, um produto, uma escala e uma soma, ao custo de varrer $C$ três vezes; a GEMM faz as três em uma única passagem.

Por que alguém quereria $\alpha$ e $\beta$? Porque a álgebra linear dos modelos está cheia deles, e cada um absorvido pela GEMM é uma varredura a menos sobre a memória. O exemplo mais direto está na atenção da Seção 13: os scores são $S=QK^\top/\sqrt{d_h}$, de modo que uma única chamada com $\operatorname{op}$ transpondo o segundo operando e $\alpha=1/\sqrt{d_h}$ calcula tudo, sem materializar $K^\top$ nem varrer $S$ de novo para dividir pela constante. Já $\beta=1$ transforma a GEMM em acumuladora: uma soma de produtos $\sum_k A_kB_k$ vira uma sequência de chamadas sobre o mesmo $C$, sem buffers temporários. Em uma biblioteca otimizada, até operações adicionais como viés e ativação podem ser incorporadas ao epílogo da GEMM, reduzindo leituras e escritas intermediárias.

A calculadora abaixo torna manipulável cada símbolo desta seção, usando as mesmas matrizes dos exemplos. Clique em um elemento de $C$ para ver o produto escalar que o gera, com os elementos de origem destacados; ajuste $\alpha$ e $\beta$ e confirme que o somatório central nunca muda. E tente ativar uma transposta no modo retangular: o erro que aparece ensina sobre dimensões internas mais do que qualquer parágrafo.

1.1 Quantos FLOPs?

Antes de contar, a atenta leitora precisa entender a unidade de contagem. Um FLOP, de floating-point operation (operação de ponto flutuante), é uma operação aritmética elementar, uma adição, subtração ou multiplicação, executada sobre números em ponto flutuante, o formato de precisão finita com que os computadores aproximam números reais e que a Seção 12 examinará por dentro. Divisões e raízes também são operações de ponto flutuante, porém mais caras e raras em GEMM; a contagem clássica as ignora. O FLOP mede quantidade de trabalho, não velocidade: a velocidade é medida em FLOPs, operações por segundo. A literatura também escreve FLOPS, com o S maiúsculo de per second, uma colisão tipográfica infeliz com o plural de FLOP. Neste artigo, FLOPs no plural é trabalho acumulado e FLOP’s é taxa.

Agora podemos contar.

A fórmula da Seção 1 diz que cada elemento $C_{ij}$ é a soma de $K$ produtos. São $K$ multiplicações e, para somar $K$ parcelas, $K-1$ adições: cada elemento custa $2K-1$ FLOPs, e a matriz inteira,

antes do epílogo. Multiplicar o somatório por $\alpha$ e somar $\beta\,C^{\text{velho}}$ acrescentam mais três operações por elemento, um total de $MN(2K+2)$. A convenção universal, no entanto, arredonda tudo isso para

e o arredondamento é honesto: a diferença relativa entre $2K-1$ e $2K$ é $1/(2K)$, e os termos do epílogo são $O(MN)$ contra o $O(MNK)$ do produto. Para o $K=4096$ de uma projeção de Transformer, a discrepância fica abaixo de um milésimo. A convenção existe para que artigos, bibliotecas e benchmarks comparem números calculados da mesma forma, e é o numerador que usaremos ao medir desempenho na Seção 15.

O fator $2$ dessa fórmula merece um parágrafo próprio, porque o hardware moderno não executa a multiplicação e a adição como instruções separadas. Em ponto flutuante, o resultado exato de uma operação em geral não cabe no formato e precisa ser arredondado. Escrevendo $\operatorname{fl}(x)$ para o número de ponto flutuante mais próximo de $x$, uma multiplicação seguida de uma adição calcula

com dois arredondamentos: o produto $a\cdot b$ é arredondado antes que a soma sequer comece. A instrução FMA, de fused multiply-add (multiplicação e adição fundidas), calcula o valor de $a\cdot b+c$ mantendo o produto internamente em largura estendida, sem arredondá-lo, e arredonda uma única vez, no final:

No modelo clássico de erro, cada arredondamento introduz um erro relativo de magnitude até $u$, o unit roundoff do formato ($u=2^{-24}$ em FP32); a FMA elimina um dos dois erros. O resultado é ao menos tão preciso quanto o da sequência separada, frequentemente melhor, e essa diferença de arredondamento terá consequências concretas na Seção 6.2.

A conexão com a GEMM é imediata. O laço mais interno do produto é uma cadeia de atualizações da forma

que é exatamente o padrão $a\cdot b+c$. Cada elemento de $C$ custa $K$ FMAs; a GEMM inteira, $MNK$ FMAs. Como cada FMA realiza uma multiplicação e uma adição, a convenção conta duas operações por instrução, e $MNK$ FMAs são os mesmos $2MNK$ FLOPs de antes: as duas contagens coincidem, como deveriam.

É dessa aritmética que nascem os picos das fichas técnicas. Um núcleo com duas unidades vetoriais de FMA, cada uma operando sobre oito float (a largura de AVX2), a $3$ GHz, entrega no máximo

por núcleo; dezesseis núcleos anunciam cerca de $1{,}5$ TFLOP/s. O número pressupõe todas as unidades FMA ocupadas em todos os ciclos, sem esperar um único byte da memória, uma condição que o restante deste artigo mostrará ser tudo menos automática.

O número de FLOPs, entretanto, não determina sozinho o tempo:

a menos que os dados cheguem às unidades de execução com velocidade suficiente. A diferença entre o pico anunciado e o desempenho observado quase sempre começa na memória.

2. O primeiro kernel: correto, simples e lento

Chegou a hora de transformar a definição em código, e a primeira decisão nem parece uma decisão: como guardar uma matriz na memória? A memória de um computador é unidimensional, um vetor gigante de bytes com endereços crescentes. Uma matriz é bidimensional. Toda representação exige, portanto, uma convenção de linearização, e a escolha dessa convenção terá consequências durante o artigo inteiro.

Adotaremos a ordem de linhas, ou row-major, a convenção de C e C++: as linhas são armazenadas uma após a outra. Em uma matriz com $N$ colunas e índices a partir de zero, antes do elemento $(i,j)$ vêm $i$ linhas completas, com $N$ elementos cada, e mais $j$ elementos da própria linha $i$. A posição linear é, portanto:

A matriz $A$ da Seção 1 torna-se, na memória, a sequência:

e o elemento $A_{1,1}=1$ ocupa a posição $1\cdot3+1=4$, como a leitora pode conferir contando a partir do zero. A convenção oposta, column-major, guarda colunas consecutivas, com posição $jM+i$; é a herança do Fortran e a convenção nativa das interfaces BLAS. Mais um motivo para o $\operatorname{op}(\cdot)$ da Seção 1 existir, já que ler uma matriz como se estivesse transposta converte uma convenção na outra sem copiar um único byte. Os fragmentos de Tensor Cores da Seção 11 usarão as duas.

Com a convenção fixada, o kernel escreve-se sozinho: um laço para $i$, um para $j$ e, dentro deles, o somatório em $p$ da definição, acumulado em um registrador antes do epílogo com $\alpha$ e $\beta$:

#include <cstddef>
#include <span>

void gemm_ijk(std::span<const float> A,
              std::span<const float> B,
              std::span<float> C,
              std::size_t M,
              std::size_t N,
              std::size_t K,
              float alpha = 1.0f,
              float beta = 0.0f) {
    for (std::size_t i = 0; i < M; ++i) {
        for (std::size_t j = 0; j < N; ++j) {
            float sum = 0.0f;
            for (std::size_t p = 0; p < K; ++p) {
                sum += A[i * K + p] * B[p * N + j];
            }
            C[i * N + j] = alpha * sum + beta * C[i * N + j];
        }
    }
}

O código segue diretamente a definição matemática. Sua correção é fácil de verificar. Seu padrão de memória, não. Para enxergá-lo, precisamos de um conceito: o passo (stride) de uma sequência de acessos é a distância, em elementos, entre dois acessos consecutivos. Passo $1$ significa varrer a memória em sequência; passo maior que isso, significa saltar.

No laço interno em $p$, os elementos de $A$ são lidos com passo $1$:

Mas os elementos de $B$ são lidos com passo $N$:

Em uma matriz row-major, percorrer uma coluna é saltar de linha em linha. Coloquemos números nesse salto. Com $N=4096$ em FP32, dois acessos consecutivos a $B$ distam $4096\times4$ bytes $=16$ KiB. Como veremos na Seção 3, a memória não entrega bytes avulsos: entrega linhas de cache de $64$ bytes, dezesseis float de uma vez. O acesso por coluna usa um desses dezesseis valores e salta para outra linha de cache: aproveitamos $6{,}25\%$ de cada transferência, e a largura de banda efetiva sobre $B$ é dividida por até dezesseis. O salto de $16$ KiB também atravessa quatro páginas de $4$ KiB a cada acesso, pressionando a TLB que a Seção 3.2 apresentará. A CPU executa os mesmos $2MNK$ FLOPs da Seção 1.1; o que muda é quantos bytes ela precisa mover para alimentá-los.

2.1 A ordem dos laços muda o tráfego

Antes de trocar a ordem dos laços, precisamos conquistar o direito de trocá-la. O kernel executa uma atualização escalar para cada tripla $(i,j,p)$:

e o valor final de cada $C_{ij}$ é a soma das suas $K$ parcelas. Em aritmética exata, a adição é comutativa e associativa: a soma não depende da ordem em que as parcelas chegam. Logo, qualquer ordem de aninhamento dos três laços, e são $3!=6$ possíveis, produz a mesma matriz $C$. Em ponto flutuante, reordenar muda os arredondamentos e, portanto, os últimos bits do resultado, uma sutileza que a Seção 6.2 tratará; nada disso torna uma ordem “errada”.

Se as seis ordens produzem o mesmo resultado, o critério de escolha é inteiramente de memória, e quem decide é o índice do laço interno, aquele que varia mais rápido:

Com o laço interno em $p$, que é o kernel que acabamos de escrever, $B$ é percorrida por colunas. Com o laço interno em $i$, duas matrizes são percorridas por colunas, o pior caso em row-major. Com o laço interno em $j$, nenhuma matriz anda por colunas: $B$ e $C$ fluem com passo $1$ e o elemento de $A$ fica parado em um registrador. É essa a ordem que queremos.

Há uma leitura algébrica elegante dessa escolha. Definamos o produto externo (outer product): dados vetores $u\in\mathbb{R}^{M}$ e $v\in\mathbb{R}^{N}$, a matriz $uv^\top\in\mathbb{R}^{M\times N}$ tem elementos $(uv^\top)_{ij}=u_iv_j$. Chamando de $a_p\in\mathbb{R}^{M}$ a coluna $p$ de $A$ e de $b_p^\top$ a linha $p$ de $B$, vale:

e a prova é elemento a elemento, em uma linha:

Com as matrizes da Seção 1, a decomposição é concreta: as colunas de $A$ são $(1,0)$, $(2,1)$ e $(0,3)$; as linhas de $B$ são $(4,1)$, $(2,0)$ e $(1,5)$; e

o mesmo resultado do produto linha-coluna da Seção 1, como o teorema exige. A ordem com $j$ interno calcula exatamente isso, uma linha de cada produto externo por vez:

Em código:

#include <algorithm>
#include <cstddef>
#include <span>

void gemm_ikj(std::span<const float> A,
              std::span<const float> B,
              std::span<float> C,
              std::size_t M,
              std::size_t N,
              std::size_t K,
              float alpha = 1.0f,
              float beta = 0.0f) {
    if (beta == 0.0f) {
        std::fill(C.begin(), C.end(), 0.0f);
    } else {
        for (float& value : C) {
            value *= beta;
        }
    }

    for (std::size_t i = 0; i < M; ++i) {
        for (std::size_t p = 0; p < K; ++p) {
            const float a = alpha * A[i * K + p];

            for (std::size_t j = 0; j < N; ++j) {
                C[i * N + j] += a * B[p * N + j];
            }
        }
    }
}

Um aparte de nomenclatura: a função chama-se gemm_ikj porque a literatura usa $k$ para o índice da soma e batizou a ordem de ikj; neste artigo o índice é $p$, para não colidir com a dimensão $K$. O laço interno executa a atualização $C_{ij}\mathrel{+}=a\cdot B_{pj}$, que as BLAS de nível 1 conhecem pelo nome AXPY ($a\cdot x+y$): um escalar parado em registrador, lido da memória uma vez a cada $N$ usos, e dois fluxos contíguos de passo $1$ que aproveitam os dezesseis float de cada linha de cache, em vez de um. A fórmula, o número de FLOPs e a complexidade $O(MNK)$ permanecem idênticos; o fator dezesseis de desperdício sobre $B$ desaparece. O ganho medido depende de $N$, das caches e do prefetcher: o laboratório da Seção 3 mostra a diferença em um modelo pequeno o bastante para ser inspecionado a olho nu, e o protocolo para medi-la em uma máquina real está na Seção 15.

A primeira lição prática é fundamental:

Complexidade assintótica descreve como o trabalho cresce. Localidade descreve quanto desse trabalho a máquina consegue alimentar.

3. A máquina não possui “uma memória”

É tentador imaginar que um processador lê qualquer posição de memória pelo mesmo preço. Uma máquina real tem uma hierarquia, e vale conhecê-la com números, ainda que aproximados:

1. registradores: algumas dezenas por núcleo, acessados no próprio ciclo da instrução; é onde o microkernel da Seção 5 manterá os acumuladores de $C$;
2. caches L1, L2 e, normalmente, uma cache compartilhada de último nível: de dezenas de KiB a dezenas de MiB, com latências típicas da ordem de 4, 12 e 40 ciclos, respectivamente;
3. memória principal (DRAM): dezenas de GiB, com latências de algumas centenas de ciclos, uma ordem de grandeza acima da cache mais lenta;
4. em uma GPU: registradores por thread, memória compartilhada por bloco, que faz o papel de uma cache administrada pelo programador, uma L2 e a memória global em HBM;
5. em sistemas distribuídos: a memória de outros aceleradores e de outros nós, atravessada por interconexões e pela rede.

A forma da pirâmide é a mesma em CPU e em GPU, e é por isso que as técnicas deste artigo migram de uma para a outra com tanta naturalidade: mudam os nomes e os tamanhos dos degraus, não a lógica. Quanto mais próxima da unidade de execução, menor tende a ser a capacidade e maior tende a ser a largura de banda. O objetivo de um kernel GEMM é trazer um bloco de dados de uma camada lenta e reutilizá-lo muitas vezes antes de descartá-lo.

3.1 Linhas de cache e localidade espacial

A CPU não busca apenas o float solicitado. Ela transfere uma linha de cache, um bloco de bytes contíguos. Se o código percorre uma linha da matriz, os próximos valores já estão próximos. Se percorre uma coluna de uma matriz row-major, cada acesso pode saltar para outra linha de cache.

Temos dois tipos principais de localidade, e cada um explora um mecanismo diferente do hardware. A localidade espacial consiste em usar endereços próximos: como a transferência mínima é a linha de 64 bytes, acessar o vizinho de um valor já carregado sai de graça, e foi exatamente essa gratuidade que o kernel da Seção 2 desperdiçou ao descer as colunas de $B$. A localidade temporal consiste em reutilizar um dado antes que ele seja expulso da cache: a capacidade é finita, a política de substituição descarta o que não é usado há mais tempo, e cada reutilização adiada é uma aposta contra o relógio.

GEMM oferece enorme reutilização potencial. Cada elemento de $A$ participa de $N$ produtos; cada elemento de $B$, de $M$ produtos. Uma implementação ingênua pode perder essa oportunidade e buscar os mesmos dados repetidamente.

3.2 TLB e páginas

Endereços virtuais precisam ser traduzidos para endereços físicos. A TLB, Translation Lookaside Buffer, guarda traduções recentes. Acessos com passos grandes podem tocar muitas páginas e pressionar a TLB, mesmo quando o volume de dados parece caber em uma cache. Por isso, uma GEMM de alto desempenho considera não apenas linhas de cache, mas também painéis contíguos, páginas e alinhamento.

A leitora não precisa acreditar na minha palavra. O laboratório abaixo executa as duas ordens de laço da Seção 2 sobre matrizes $8\times8$, com linhas de cache de quatro elementos e uma cache de oito linhas com política LRU, tudo pequeno o bastante para ser visto a olho nu. Execute uma ordem, anote a taxa de acerto, execute a outra. A álgebra é a mesma; o tráfego, não.

4. Blocagem para reutilizar

Dividimos as matrizes em blocos:

nos quais $I$, $J$ e $P$ representam intervalos de linhas e colunas. Em vez de consumir uma linha ou coluna inteira, processamos pequenos painéis que cabem em níveis rápidos da memória.

Se escolhermos blocos $B_M\times B_K$, $B_K\times B_N$ e $B_M\times B_N$, o conjunto de trabalho aproximado em bytes será:

em que $s$ é o número de bytes por elemento. Os tamanhos devem deixar espaço para outros dados e associatividade da cache; portanto, não basta igualar $Q_{\text{tile}}$ à capacidade nominal.

#include <algorithm>
#include <cstddef>
#include <span>

void gemm_blocked(std::span<const float> A,
                  std::span<const float> B,
                  std::span<float> C,
                  std::size_t M,
                  std::size_t N,
                  std::size_t K,
                  std::size_t BM = 64,
                  std::size_t BN = 64,
                  std::size_t BK = 64) {
    std::fill(C.begin(), C.end(), 0.0f);

    for (std::size_t ii = 0; ii < M; ii += BM) {
        const std::size_t i_end = std::min(ii + BM, M);

        for (std::size_t jj = 0; jj < N; jj += BN) {
            const std::size_t j_end = std::min(jj + BN, N);

            for (std::size_t pp = 0; pp < K; pp += BK) {
                const std::size_t p_end = std::min(pp + BK, K);

                for (std::size_t i = ii; i < i_end; ++i) {
                    for (std::size_t p = pp; p < p_end; ++p) {
                        const float a = A[i * K + p];

                        for (std::size_t j = jj; j < j_end; ++j) {
                            C[i * N + j] += a * B[p * N + j];
                        }
                    }
                }
            }
        }
    }
}

Esse código é didático, não uma substituição para BLAS. Ainda assim, ele contém a ideia central: um bloco de $C$ é atualizado por blocos compatíveis de $A$ e $B$, mantendo dados quentes na hierarquia.

4.1 Um tamanho de bloco não serve para tudo

O código anterior usa BM = BN = BK = 64 com a confiança serena de quem copiou o valor de um tutorial. Não há nada de sagrado nesses números. A escolha dos blocos é um problema de alocação de recursos escassos, e cada nível da hierarquia impõe a sua própria restrição. Vale a pena percorrer essas restrições uma a uma, porque esse é exatamente o raciocínio que as bibliotecas BLAS fazem por nós, e que precisaremos refazer sempre que uma GEMM não entregar o desempenho esperado.

Comece pelo nível mais rápido. O microkernel, que veremos em detalhe na Seção 5, mantém um bloco $M_R\times N_R$ de $C$ em registradores durante todo o laço em $K$. O tamanho desse bloco é ditado pela quantidade de registradores vetoriais disponíveis. Em AVX2 há $16$ registradores de $256$ bits, cada um comportando oito float. Um microkernel $6\times16$ mantém $6\times16=96$ acumuladores de $C$ em $12$ registradores, seis linhas com dois vetores de oito colunas cada, sobrando quatro registradores para carregar elementos de $A$ e $B$ e alimentar as FMAs. Não por acaso, $6\times16$ é o formato do microkernel SGEMM que BLIS e OpenBLAS usam em processadores da classe Haswell. Com AVX-512, que oferece $32$ registradores de $512$ bits, o microkernel cresce; com NEON em ARMv8, que oferece $32$ registradores de apenas $128$ bits, ele assume outra forma:

A tabela também explica por que o tipo numérico altera a geometria. Em double, cada registrador comporta metade dos elementos e o microkernel encolhe: o mesmo Haswell usa $6\times8$ para DGEMM. Em FP16 nativo, dobra-se a quantidade de elementos por registrador e a forma muda de novo. Mudar o tipo não muda apenas a precisão: muda o formato de todos os blocos construídos em volta do microkernel.

Subindo um nível, as caches ditam $B_M$, $B_K$ e $B_N$. A decomposição clássica de Goto e van de Geijn, que as bibliotecas modernas seguem com pequenas variações, atribui cada painel a um nível específico da hierarquia:

1. o painel estreito $B_K\times N_R$ de $B$ que o microkernel consome deve residir na L1;
2. o painel $B_M\times B_K$ de $A$, empacotado, deve residir na L2;
3. o painel $B_K\times B_N$ de $B$, empacotado, deve residir na L3.

Com números concretos o critério deixa de ser abstrato. Um Haswell tem L1 de dados de $32$ KiB e L2 de $256$ KiB por núcleo. Escolhendo $B_K=256$, o painel estreito de $B$ ocupa $256\times16\times4$ bytes $=16$ KiB, metade da L1. Os parâmetros do BLIS para SGEMM nessa arquitetura são da ordem de $B_M=144$, $B_K=256$ e $B_N=4080$: o painel de $A$ ocupa $144\times256\times4\approx144$ KiB, pouco mais da metade da L2, e o painel de $B$ ocupa cerca de $4$ MiB, confortável na L3. Nenhum desses valores é $64$, e nenhum é potência de dois.

A atenta leitora deve ter estranhado o “metade”. Por que não usar a capacidade inteira da cache? Porque a cache não guarda apenas o nosso painel. As linhas de $C$ que o microkernel lê e escreve passam pelo mesmo caminho, o prefetcher traz dados adiantados, a pilha e o código também ocupam espaço, e caches associativas por conjuntos sofrem conflitos: endereços separados por passos que são potências de dois podem disputar os mesmos conjuntos e expulsar uns aos outros mesmo com capacidade nominal de sobra. Dimensionar um bloco para $100\%$ da cache é planejar uma mudança contando cada centímetro cúbico do caminhão, sem deixar espaço para fechar a porta.

As dimensões das matrizes impõem a última restrição, e é a que mais aparece em Transformers. Blocos de $144\times256$ pressupõem que existam pelo menos $144$ linhas para bloquear. No decode de um modelo de linguagem, como veremos na Seção 14, a projeção de um único token tem $M=1$: não há o que bloquear no eixo $M$, e o kernel ideal é outro, mais próximo de um GEMV cuidadosamente vetorizado. Matrizes com $K$ pequeno, como as GEMMs por cabeça de atenção em que $K=d_h$ vale $64$ ou $128$, limitam a reutilização ao longo do laço interno e favorecem blocos achatados. Por isso bibliotecas como cuBLASLt aceitam descrições detalhadas do problema e escolhem entre dezenas de kernels, e projetos como o libxsmm geram kernels sob medida para matrizes pequenas em tempo de execução.

Threads, por fim, dividem os mesmos recursos: a L3 é compartilhada, de modo que $B_N$ efetivo por thread diminui conforme os núcleos aumentam, e cada thread precisa de seus próprios buffers de packing. O bloco ótimo para um núcleo raramente é o bloco ótimo para trinta e dois.

Diante de tantas variáveis, as bibliotecas adotam duas estratégias. O ATLAS media exaustivamente combinações durante a instalação, uma forma honesta de admitir ignorância. O BLIS mostrou que um modelo analítico da hierarquia, alimentado com capacidade, associatividade e largura de linha de cada cache, chega perto do ótimo sem busca exaustiva. Para o nosso código didático, a lição é mais modesta: em vez de assumir que 64 é mágico, medimos vários tamanhos em matrizes representativas da carga real, na máquina real.

4.2 Blocos de borda

Dimensões reais raramente são múltiplas perfeitas do bloco. O código anterior resolve o problema com std::min, o que é correto e esconde um custo fácil de subestimar. Tome $M=N=K=1000$ com blocos de $64$. Como $1000=15\cdot64+40$, sobra uma borda de $40$ em cada eixo, e a fração dos elementos de $C$ que pertence a algum bloco parcial nos eixos $M$ e $N$ é

Parece desprezível. Mas suponha que essa borda seja processada por um laço escalar oito vezes mais lento que o caminho vetorizado em AVX2. O tempo total fica em torno de $0{,}922\cdot1+0{,}078\cdot8\approx1{,}55$ vez o tempo do caso perfeito: menos de $8\%$ do trabalho consumiu mais de um terço do tempo. É a lei de Amdahl em miniatura, com a agravante de que a borda no eixo $K$ ainda adiciona sua própria parcela. Kernels industriais atacam o problema com quatro estratégias, em ordem crescente de sofisticação.

A primeira é o laço escalar de cauda: o caminho vetorizado processa múltiplos inteiros da largura do vetor e um laço simples termina o resto. O kernel AVX2 da Seção 6 faz exatamente isso no segundo laço em j. A estratégia é portátil, trivialmente correta e aceitável quando as matrizes são grandes o bastante para diluir a borda, o que a nossa conta acima mostra que nem sempre é o caso.

A segunda estratégia usa máscaras vetoriais. AVX-512 introduziu registradores de máscara que desligam lanes individuais de uma instrução:

const unsigned remainder = static_cast<unsigned>(N - j);
const __mmask16 mask =
    static_cast<__mmask16>((1u << remainder) - 1u);

__m512 acc = _mm512_maskz_loadu_ps(mask, C + i * N + j);
// FMAs usando cargas mascaradas de B...
_mm512_mask_storeu_ps(C + i * N + j, mask, acc);

Lanes desligadas não carregam, não computam e, sobretudo, não escrevem: a borda executa no mesmo caminho vetorizado, sem tocar memória fora da matriz. O SVE do ARM leva a ideia ao extremo com predicados gerados por instruções como whilelt, tornando cada iteração naturalmente mascarada; o conceito de cauda simplesmente desaparece do código. AVX2 fica no meio do caminho, com _mm256_maskload_ps e _mm256_maskstore_ps, funcionais porém mais caros que suas contrapartes plenas.

A terceira estratégia aproveita o packing da Seção 5. Se os painéis de $A$ e $B$ já serão copiados para buffers contíguos, custa pouco arredondar suas dimensões para múltiplos de $M_R$ e $N_R$ e preencher a diferença com zeros. Zeros atravessam FMAs sem alterar o resultado, pois $c+a\cdot0=c$, e o microkernel passa a executar sempre no formato pleno, sem casos especiais. Dois cuidados: $C$ nunca é preenchido, então as escritas do microkernel precisam ser restritas à região real, com máscaras ou variantes de armazenamento parcial; e os FLOPs fantasmas têm custo, limitado a $M_R-1$ linhas e $N_R-1$ colunas extras por painel, irrelevante para matrizes grandes e mensurável para pequenas.

A quarta estratégia é manter uma família de microkernels: o formato pleno $M_R\times N_R$ e variantes para cada largura parcial, selecionadas em tempo de execução. O caso extremo é o libxsmm, que gera via JIT um kernel específico para as dimensões exatas do problema, abordagem que nasceu precisamente porque cargas de inferência multiplicam matrizes pequenas, nas quais a borda não é exceção, é o problema inteiro. O preço é tamanho de código e complexidade de despacho.

Na GPU, o problema muda de nome, não de natureza. O kernel CUDA da Seção 10 trata bordas por predicação: as cargas usam o operador ternário para preencher a memória compartilhada com zeros e o if (row < M && col < N) final restringe as escritas. O custo mais sutil é a quantização de tiles: a grade tem $\lceil M/T\rceil\times\lceil N/T\rceil$ blocos, e blocos são lançados inteiros. Com $M=N=65$ e $T=32$, lançamos $3\times3=9$ blocos, ou seja, $9\cdot1024=9216$ threads para $65^2=4225$ elementos úteis: menos da metade do trabalho lançado produz resultado. Tensor Cores agravam a exigência, porque operam sobre fragmentos de forma fixa, como o $16\times16\times16$ da Seção 11, e o cuBLAS obtém seus kernels mais rápidos quando dimensões e alinhamentos respeitam múltiplos de $8$ ou $16$ elementos em FP16.

Isso explica um hábito da arquitetura de modelos que, à primeira vista, parece numerologia: $d_{\text{model}}$ de $4096$ no Llama, $12288$ no GPT-3, cabeças de dimensão $64$ ou $128$, vocabulários arredondados para múltiplos de $64$ ou $128$, como faz o Megatron-LM ao preencher a matriz de embeddings. Quando toda dimensão de toda GEMM é múltipla dos tiles do hardware, não existem blocos de borda, não existe quantização de tiles e todos os caminhos rápidos ficam disponíveis. Quando um artigo descreve um modelo com dimensões suspeitosamente redondas, há um engenheiro de kernels em algum lugar respirando aliviado.

O laboratório seguinte permite experimentar o compromisso central desta seção. Escolha a dimensão da matriz, o tamanho do bloco e o tipo numérico, e observe as duas consequências simultâneas: o conjunto de trabalho contra as caches e o tráfego estimado até a DRAM. Note que aumentar o bloco sempre reduz o tráfego no modelo, até o momento em que ele deixa de caber, e é nesse conflito que vive toda a engenharia da Seção 4.1.

5. Packing, macro-kernel e microkernel

A blocagem ainda lê painéis na disposição original, que pode possuir passos inconvenientes. Bibliotecas como BLIS e OpenBLAS e implementações derivadas das ideias de Goto usam packing: copiam painéis de $A$ e $B$ para buffers contíguos em uma disposição adequada ao microkernel.

A cópia parece trabalho extra. Ela compensa quando o painel empacotado é reutilizado vezes suficientes.

Uma GEMM de alto desempenho pode ser pensada em três níveis:

1. packing: copia painéis de $A$ e $B$ para buffers contíguos e alinhados, no layout exato que o microkernel consome; paga-se um custo de cópia $O(n^2)$ para comprar acessos de passo $1$ dentro do laço $O(n^3)$, uma troca que a aritmética favorece;
2. macro-kernel: percorre os blocos dimensionados para cada nível de cache, com os parâmetros $B_M$, $B_K$ e $B_N$ que a Seção 4.1 derivou;
3. microkernel: calcula um pequeno bloco $M_R\times N_R$ de $C$, mantendo seus acumuladores em registradores do início ao fim do laço em $K$.

Guarde essa estrutura de três níveis: ela reaparecerá inteira na GPU da Seção 10, com a memória compartilhada no papel da cache que o programador controla e microtiles por warp no papel do microkernel. A mesma ideia, em outra hierarquia.

O microkernel idealiza:

Durante o laço em $K_C$, os $M_RN_R$ valores parciais de $C$ permanecem em registradores. Assim, não lemos e escrevemos $C$ a cada multiplicação. Carregamos no início e armazenamos no fim.

5.1 Por que registradores importam?

Considere uma atualização escalar:

Se $C_{ij}$ for lido e escrito na cache para cada $p$, o tráfego cresce com $K$. Se permanecer em um registrador durante todo o painel, há uma leitura e uma escrita por bloco, enquanto muitas FMAs acontecem internamente.

O tamanho do microkernel é um compromisso. Poucos acumuladores desperdiçam unidades vetoriais; acumuladores demais causam register spilling, isto é, o compilador derrama valores para a memória, destruindo a vantagem.

6. SIMD: várias colunas por instrução

SIMD, Single Instruction, Multiple Data, aplica uma instrução a vários elementos. Em AVX2, um registrador de 256 bits comporta oito valores float. Uma FMA vetorial pode atualizar oito elementos de $C$ simultaneamente:

O kernel seguinte ilustra um acumulador $1\times 8$. Ele é propositalmente menor que um microkernel industrial, mas conserva $C$ nos registradores durante o laço em $K$.

#include <cstddef>
#include <immintrin.h>

void gemm_avx2_1x8(const float* A,
                   const float* B,
                   float* C,
                   std::size_t M,
                   std::size_t N,
                   std::size_t K,
                   float alpha = 1.0f,
                   float beta = 0.0f) {
    for (std::size_t i = 0; i < M; ++i) {
        std::size_t j = 0;

        for (; j + 8 <= N; j += 8) {
            __m256 acc = _mm256_loadu_ps(C + i * N + j);
            acc = _mm256_mul_ps(acc, _mm256_set1_ps(beta));

            for (std::size_t p = 0; p < K; ++p) {
                const __m256 a =
                    _mm256_set1_ps(alpha * A[i * K + p]);
                const __m256 b =
                    _mm256_loadu_ps(B + p * N + j);

                acc = _mm256_fmadd_ps(a, b, acc);
            }

            _mm256_storeu_ps(C + i * N + j, acc);
        }

        for (; j < N; ++j) {
            float acc = beta * C[i * N + j];

            for (std::size_t p = 0; p < K; ++p) {
                acc += alpha * A[i * K + p] * B[p * N + j];
            }

            C[i * N + j] = acc;
        }
    }
}

Em GCC ou Clang, uma compilação típica usaria -O3 -mavx2 -mfma. O executável só pode chamar esse kernel em processadores que ofereçam essas extensões. Em software distribuído, precisamos de detecção de recursos e versões alternativas.

A implementação usa loadu, que aceita endereços não alinhados. Buffers alinhados e load alinhado podem ajudar em alguns cenários, mas alinhamento sozinho não substitui localidade, blocagem e reutilização.

6.1 Autovetorização

Antes de escrever intrínsecos, convém dar ao compilador a chance de fazer o trabalho sozinho. GCC, Clang e MSVC vetorizam laços automaticamente há muitos anos, e o kernel gemm_ikj da Seção 2.1 é um candidato quase perfeito: o laço interno percorre $B$ e $C$ contiguamente, multiplicando por um escalar reutilizado. Quase perfeito. Os obstáculos que restam são instrutivos, porque são os mesmos que aparecem em qualquer código numérico.

O primeiro é o aliasing. Para vetorizar C[i*N+j] += a * B[p*N+j], o compilador precisa provar que escrever em $C$ não altera os valores de $B$ que serão lidos nas próximas iterações. Olhando apenas para a assinatura da função, ele não pode provar nada: os intervalos de memória poderiam se sobrepor, e std::span não carrega nenhuma promessa em contrário. Diante da dúvida, o compilador escolhe entre gerar duas versões do laço com um teste de sobreposição em tempo de execução, o caso mais comum, ou desistir da vetorização. Podemos eliminar a dúvida com a qualificação de não aliasing: restrict em C, ou a extensão __restrict, aceita pelos três grandes compiladores de C++:

void gemm_ikj(const float* __restrict A,
              const float* __restrict B,
              float* __restrict C,
              std::size_t M, std::size_t N, std::size_t K);

A promessa passa a ser nossa responsabilidade. Se os ponteiros se sobrepuserem, o comportamento é indefinido, e o compilador, tendo confiado em nós, não emitirá nenhum aviso.

O segundo obstáculo é a redução, e ele explica por que a ordem ijk original resiste à vetorização mesmo sem aliasing. O laço interno de gemm_ijk acumula sum += A[i*K+p] * B[p*N+j]. Além do acesso a $B$ com passo $N$, vetorizar essa soma exige quebrá-la em acumuladores parciais e reassociar as adições no final. Em aritmética de ponto flutuante, reassociar muda o resultado, como discutiremos na Seção 6.2, e o compilador não faz isso sem permissão explícita. As permissões possíveis são -ffast-math ou -fassociative-math, que liberam a reassociação para o programa inteiro com todas as consequências numéricas, ou, de forma mais cirúrgica, uma anotação no laço específico:

#pragma omp simd reduction(+:sum)
for (std::size_t p = 0; p < K; ++p) {
    sum += A[i * K + p] * B[p * N + j];
}

Com -fopenmp-simd, o pragma é honrado sem vincular o programa ao runtime completo de OpenMP. A anotação declara que aceitamos a reassociação naquele ponto, e apenas naquele ponto.

O terceiro obstáculo é a estrutura do laço. O vetorizador quer limites que não mudem durante a execução, ausência de saídas antecipadas e ausência de chamadas de função opacas no corpo. Nossos kernels já satisfazem essas condições; código real nem sempre.

Resolvidos os obstáculos, resta saber se o compilador de fato vetorizou, e aqui entra a ferramenta mais subutilizada da otimização: os relatórios de vetorização. Em Clang, -Rpass=loop-vectorize informa os laços vetorizados e -Rpass-missed=loop-vectorize, com -Rpass-analysis=loop-vectorize, informa os recusados e o motivo. Em GCC, -fopt-info-vec-missed cumpre o mesmo papel; em MSVC, /Qvec-report:2. Uma mensagem como “loop not vectorized: cannot prove it is safe to reorder floating-point operations” aponta diretamente para a redução; uma menção a “runtime pointer checks” aponta para o aliasing. Ler o relatório transforma adivinhação em diagnóstico. Convém também lembrar das flags de arquitetura: sem -mavx2 -mfma ou -march=native, um compilador para x86-64 assume a linha de base SSE2 e vetoriza com registradores de $128$ bits, entregando metade ou um quarto da largura disponível na máquina.

Há, porém, um teto para tudo isso. A autovetorização opera sobre o laço que escrevemos: ela não inventa blocagem, não faz packing e não reestrutura três laços aninhados em um microkernel $6\times16$ com acumuladores em registradores. O compilador vetoriza a aritmética; a coreografia de dados continua sendo nossa. Por isso a divisão de trabalho madura é usar autovetorização com __restrict e relatórios em todo o código morno, e reservar os intrínsecos para o microkernel, onde cada registrador importa e o custo de manutenção se paga.

6.2 FMA e reprodutibilidade

FMA calcula $ab+c$ com um único arredondamento. Separar multiplicação e adição produz dois. Por isso, uma versão vetorizada pode diferir nos últimos bits da versão escalar. A ordem das somas também muda com blocagem, paralelismo e redução.

Isso não significa que uma das implementações esteja necessariamente errada. Testes numéricos devem usar tolerâncias relacionadas ao tipo, a $K$ e à magnitude dos operandos, além de referências em precisão maior quando a aplicação exigir.

7. Paralelismo em CPUs

Blocos distintos de $C$ podem ser calculados por threads distintas, pois as escritas não se sobrepõem. Uma decomposição comum distribui blocos nos eixos $M$ e $N$. Entretanto, adicionar threads não multiplica recursos: divide alguns e cria custos novos.

O recurso dividido mais importante é a largura de banda da memória: dezesseis núcleos calculando não recebem dezesseis vezes mais bytes por segundo, e um kernel que já era limitado por memória com um núcleo continua limitado com dezesseis, apenas com mais elegância. A cache de último nível também é compartilhada, de modo que o bloco efetivo disponível por thread encolhe conforme os núcleos aumentam, como a Seção 4.1 observou. Os custos novos começam na criação e na sincronização das threads, que só se amortizam quando há trabalho suficiente por thread, e culminam no falso compartilhamento (false sharing): duas threads escrevendo em elementos distintos que caem na mesma linha de cache de 64 bytes obrigam o protocolo de coerência a rebater a linha entre os núcleos, serializando escritas que eram logicamente independentes — um risco real nas bordas entre blocos vizinhos de $C$. Por fim, a afinidade entre threads e núcleos importa: o sistema operacional pode migrar threads, e cada migração esfria as caches que o kernel tinha aquecido.

Para matrizes grandes, bibliotecas BLAS paralelas já resolvem boa parte desse problema. Criar paralelismo externo sobre uma BLAS que também cria threads pode provocar oversubscription. O resultado são mais threads que núcleos úteis, mais trocas de contexto e menor desempenho.

Em máquinas NUMA, páginas de memória possuem afinidade com nós. Inicializar uma matriz em uma única thread e depois distribuí-la entre vários sockets pode forçar acessos remotos. Política de alocação, first touch e afinidade passam a fazer parte do algoritmo observado.

8. Custo de movimentação de dados e intensidade aritmética

Definimos a intensidade aritmética, ou operacional, como:

na qual $W$ é o trabalho em FLOPs e $Q$ é o tráfego de bytes entre o kernel e um nível de memória.

Para GEMM com $\beta=0$, um limite inferior idealizado para o tráfego é ler $A$, ler $B$ e escrever $C$ uma vez:

Se $\beta\neq0$, também precisamos ler o valor anterior de $C$ antes de escrever o novo:

Logo,

Para matrizes quadradas $N\times N$:

Em FP32, $s=4$, portanto $I_{\max}\approx N/6$ FLOP/byte. Em FP16 ou BF16, ignorando conversões e acumuladores, o denominador cai para dois bytes por elemento e a intensidade potencial aumenta.

Esse é um limite idealizado. Se os mesmos blocos forem recarregados da memória lenta, $Q$ será maior e a intensidade efetiva cairá. Packing, blocagem e fusão existem justamente para aproximar o kernel desse limite.

8.1 Um exemplo numérico

Para uma GEMM quadrada $4096\times4096$ em FP16, com $\beta=0$:

O tráfego mínimo idealizado é:

A intensidade máxima idealizada fica próxima de:

Há trabalho suficiente para reutilizar dados e manter unidades aritméticas ocupadas. Compare isso com uma multiplicação matriz-vetor em que cada peso é usado uma única vez: a intensidade pode ficar próxima de apenas uma operação útil por byte de pesos. Essa diferença será decisiva no prefill e na decodificação de modelos de linguagem.

9. O modelo Roofline

O modelo Roofline relaciona desempenho, intensidade aritmética, largura de banda e pico computacional:

Temos dois regimes:

• limitado por memória: $B_{\text{memória}}I<P_{\text{pico}}$;
• limitado por computação: $B_{\text{memória}}I\geq P_{\text{pico}}$.

O ponto de transição é:

Se o kernel possui $I<I^\star$, aumentar a quantidade de unidades aritméticas não resolve o gargalo. Precisamos reduzir tráfego, aumentar reutilização ou usar menos bytes por elemento. Se $I>I^\star$, então vetorização, Tensor Cores, ocupação e emissão de instruções ganham importância.

O Roofline não prevê todos os detalhes: latência, dependências, instruções especiais, ocupação e comunicação também importam. Sua força está em impedir diagnósticos equivocados. Antes de “otimizar FLOPs”, perguntamos em qual teto o kernel bate.

Como o Roofline é um modelo gráfico, nada mais justo que um gráfico interativo. Ajuste o pico e a largura de banda, ou use os presets de hardware, e observe onde caem as duas operações típicas de um Transformer em inferência: a GEMM quadrada do prefill e o GEMV do decode, que a Seção 14 discutirá em detalhe. Aumente o lote do decode e veja o ponto deslizar ao longo do teto inclinado: é o argumento inteiro do batching em uma imagem.

10. Da CPU para a GPU

Uma GPU executa muitas threads e depende de paralelismo massivo para esconder latências. Uma GEMM simples pode atribuir um elemento de $C$ a cada thread. Isso é correto, mas cada thread volta à memória global para buscar dados que threads vizinhas também usam.

A solução repete a ideia da CPU em outra hierarquia: tiling.

1. um bloco de threads carrega um tile de $A$ e um tile de $B$;
2. os tiles ficam em memória compartilhada, no chip;
3. threads reutilizam esses dados para acumular um tile de $C$;
4. o processo avança pelo eixo $K$.

template <int TILE>
__global__ void gemm_tiled_cuda(const float* __restrict__ A,
                                const float* __restrict__ B,
                                float* __restrict__ C,
                                int M,
                                int N,
                                int K) {
    __shared__ float tile_a[TILE][TILE];
    __shared__ float tile_b[TILE][TILE];

    const int row = blockIdx.y * TILE + threadIdx.y;
    const int col = blockIdx.x * TILE + threadIdx.x;

    float acc = 0.0f;
    const int tiles_k = (K + TILE - 1) / TILE;

    for (int tile = 0; tile < tiles_k; ++tile) {
        const int a_col = tile * TILE + threadIdx.x;
        const int b_row = tile * TILE + threadIdx.y;

        tile_a[threadIdx.y][threadIdx.x] =
            (row < M && a_col < K) ? A[row * K + a_col] : 0.0f;

        tile_b[threadIdx.y][threadIdx.x] =
            (b_row < K && col < N) ? B[b_row * N + col] : 0.0f;

        __syncthreads();

        #pragma unroll
        for (int p = 0; p < TILE; ++p) {
            acc += tile_a[threadIdx.y][p]
                 * tile_b[p][threadIdx.x];
        }

        __syncthreads();
    }

    if (row < M && col < N) {
        C[row * N + col] = acc;
    }
}

Esse kernel serve para explicar a estrutura. cuBLAS, CUTLASS e kernels de compiladores especializados usam blocagem em múltiplos níveis, cópias assíncronas, pipelines, microtiles por thread ou warp e instruções matriciais.

10.1 Acessos coalescidos

Threads de um warp devem acessar endereços que possam ser combinados em poucas transações de memória. No carregamento anterior, threadIdx.x varia ao longo de colunas contíguas. Isso favorece acessos coalescidos.

A memória compartilhada também é dividida em bancos. Padrões ruins podem criar conflitos e serializar acessos. Às vezes uma coluna extra de padding altera o passo e remove conflitos:

__shared__ float transposed[TILE][TILE + 1];

O +1 não muda a matemática; muda o mapeamento dos endereços nos bancos.

10.2 Ocupação não é um objetivo isolado

A ocupação mede quantos warps ativos podem residir em um multiprocessador. Blocos maiores podem aumentar o paralelismo, mas também consumir mais registradores e memória compartilhada. Se o consumo por bloco for alto, menos blocos residirão simultaneamente.

O melhor kernel não é necessariamente o de maior ocupação. Um kernel com mais dados reutilizados e mais trabalho por thread pode vencer com ocupação menor. Medimos desempenho, tráfego, utilização das unidades e stall reasons em conjunto.

11. Tensor Cores e operações MMA

Unidades SIMD aplicam a mesma operação a vários escalares. Tensor Cores executam operações matriciais do tipo matrix multiply-accumulate sobre pequenos tiles:

Eles não multiplicam magicamente uma matriz inteira. O software continua responsável por toda a coreografia em volta: decompor as matrizes em tiles do formato exato que a unidade aceita, transportar os operandos da memória global até os registradores passando pela memória compartilhada, escolher layouts e tipos compatíveis com o hardware, alimentar a unidade em ritmo constante para que ela não passe ciclos ociosa, acumular e armazenar os resultados, e tratar as bordas da Seção 4.2, agora agravadas por formas fixas de tile. Em outras palavras, os Tensor Cores substituem o miolo aritmético do microkernel da Seção 5; todo o resto daquela estrutura continua sendo problema nosso.

Na interface WMMA de CUDA, um warp coopera sobre fragmentos matriciais. O exemplo conceitual seguinte usa operandos de meia precisão e acumuladores FP32 em um tile suportado por determinadas arquiteturas:

#include <cuda_fp16.h>
#include <mma.h>

using namespace nvcuda;

__device__ void mma_tile(const half* a,
                         const half* b,
                         float* c,
                         int lda,
                         int ldb,
                         int ldc) {
    wmma::fragment<wmma::matrix_a,
                   16, 16, 16,
                   half,
                   wmma::row_major> a_frag;

    wmma::fragment<wmma::matrix_b,
                   16, 16, 16,
                   half,
                   wmma::col_major> b_frag;

    wmma::fragment<wmma::accumulator,
                   16, 16, 16,
                   float> c_frag;

    wmma::fill_fragment(c_frag, 0.0f);
    wmma::load_matrix_sync(a_frag, a, lda);
    wmma::load_matrix_sync(b_frag, b, ldb);
    wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);
wmma::store_matrix_sync(c, c_frag, ldc,
                            wmma::mem_row_major);
}

O fragmento $A$ está em ordem de linhas e o fragmento $B$, em ordem de colunas; os ponteiros e leading dimensions precisam respeitar esses contratos. As formas, os tipos e os requisitos variam por arquitetura. Em aplicações, normalmente chamamos cuBLAS ou cuBLASLt e deixamos a biblioteca escolher kernels compatíveis. Dimensões e alinhamentos adequados continuam importantes para atingir o melhor desempenho, mesmo quando a biblioteca consegue tratar casos gerais.

11.1 Tensor Cores não eliminam o gargalo de memória

Uma unidade matricial pode consumir operandos muito rapidamente. Se os tiles não forem trazidos e reutilizados no ritmo necessário, ela ficará ociosa. E quanto maior o pico aritmético, mais severa a obrigação de alimentá-lo: é o Roofline da Seção 9 com o teto horizontal mais alto e o ponto de equilíbrio $I^\star$ empurrado para a direita. Na prática, isso significa usar memória compartilhada e registradores para que cada byte trazido seja reutilizado o máximo possível, sobrepor cópia e computação para esconder a latência da memória atrás do trabalho útil, manter os acessos coalescidos da Seção 10.1, empregar tiles grandes o bastante para diluir o custo fixo de cada transferência, e fundir epílogos para não materializar intermediários. Nenhum item dessa receita é novo; todos ficaram mais urgentes.

Tensor Cores ampliam o teto horizontal do Roofline. A inclinação do teto de memória continua existindo.

12. Precisão mista

Reduzir a precisão ataca o gargalo deste artigo pelos dois lados. Do lado da memória, metade dos bits significa metade dos bytes transferidos: em FP16, a mesma linha de cache de 64 bytes carrega 32 valores em vez de 16, e a intensidade aritmética máxima da Seção 8 dobra sem que uma única linha de algoritmo mude. Do lado da computação, registradores e transações passam a conter o dobro de elementos, e os formatos reduzidos habilitam as unidades matriciais da Seção 11, cujos picos em FP16 superam em muito os de FP32. Há ainda o efeito de capacidade: pesos, ativações e gradientes ocupando menos memória significam modelos maiores, lotes maiores ou contextos mais longos no mesmo hardware.

Mas “usar menos bits” não é uma transformação semanticamente gratuita.

A tabela omite detalhes como valores especiais, subnormais e variantes de implementação. Nem todo hardware oferece o mesmo conjunto de formatos.

12.1 Armazenamento, produto e acumulação podem ter tipos diferentes

Uma GEMM de precisão mista pode armazenar $A$ e $B$ em FP16 ou BF16, multiplicá-los em baixa precisão e acumular os resultados em FP32:

Isso importa porque cada elemento de $C$ soma $K$ produtos. Acumular em baixa precisão pode perder contribuições pequenas à medida que o acumulador cresce. FP32 não recupera bits já perdidos na representação dos operandos, mas protege melhor a redução.

Precisamos distinguir pelo menos quatro escolhas:

1. tipo de armazenamento de $A$ e $B$;
2. precisão efetiva do produto;
3. tipo do acumulador;
4. tipo em que $C$ é armazenado.

Dizer apenas “a operação usa FP16” é insuficiente.

12.2 Erro de uma soma longa

No modelo clássico de arredondamento, uma soma de $K$ termos em precisão com erro unitário $u$ pode acumular um fator aproximado:

Esse limite é conservador, mas mostra por que o comprimento da redução importa. Reordenar a soma, usar FMA e acumular em precisão maior muda o comportamento numérico.

Para validar uma GEMM de baixa precisão, não basta comparar bit a bit com FP32: os bits vão diferir, e a pergunta certa é se diferem dentro do esperado. Avaliamos o erro absoluto e o relativo elemento a elemento, contra as tolerâncias da Seção 15.3; a norma do erro, que resume a matriz inteira em um número comparável com o limite $\gamma_K$; a distribuição dos valores, porque underflow e saturação aparecem nas caudas muito antes de aparecerem na média; e, em aprendizado de máquina, o critério final é funcional: o impacto na perda, nas métricas do modelo e na estabilidade do treinamento ao longo de muitos passos, já que um viés minúsculo por passo pode compor em divergência.

12.3 Treinamento e loss scaling

Gradientes pequenos podem sofrer underflow em FP16. O loss scaling multiplica a perda por um fator $S$ antes da retropropagação:

Os gradientes calculados também são escalados:

Antes da atualização, dividimos por $S$. Escalonamento dinâmico reduz $S$ quando aparecem Inf ou NaN e pode aumentá-lo após passos estáveis. BF16 possui faixa de expoentes maior que FP16 e frequentemente reduz a necessidade desse mecanismo, mas não elimina toda fonte de instabilidade.

Em treinamento de precisão mista também é comum preservar estados do otimizador e, dependendo do método, uma cópia dos pesos em FP32.

12.4 Quantização não é sinônimo de precisão mista

Precisão mista escolhe tipos diferentes para armazenamento e computação. Quantização introduz uma escala que mapeia valores reais para um conjunto discreto, frequentemente inteiro:

na qual $s$ é a escala e $z$ é o ponto zero. INT8, INT4 e formatos em blocos podem reduzir drasticamente tráfego, mas exigem escalas, calibração ou treinamento consciente de quantização. O gargalo pode migrar para conversões, desquantização e kernels de formas específicas.

13. GEMM dentro de um Transformer

Agora conectaremos a engenharia à arquitetura.

Considere uma entrada:

com lote $B$, comprimento $L$ e dimensão do modelo $d$. Para projeções lineares, podemos achatar os dois primeiros eixos:

13.1 Projeções Q, K e V

As projeções são:

com $W_Q,W_K,W_V\in\mathbb{R}^{d\times d}$. Podemos concatenar os pesos:

e executar uma única GEMM:

A fusão reduz lançamentos de kernels e permite ler $X$ uma vez para produzir as três projeções. O trabalho aproximado é:

A projeção de saída acrescenta aproximadamente $2BLd^2$ FLOPs.

13.2 Atenção como duas multiplicações em lote

Para cada cabeça:

$QK^T$ e $PV$ são multiplicações matriciais em lote. Juntas, custam aproximadamente:

Entretanto, a softmax separa as duas GEMMs. Uma implementação ingênua:

1. grava $S$ na HBM;
2. lê $S$ para aplicar máscara e softmax;
3. grava $P$;
4. lê $P$ para multiplicar por $V$.

A matriz intermediária possui dimensão $B\times h\times L\times L$. O custo quadrático de armazenamento e tráfego pode dominar antes mesmo de os FLOPs serem o problema principal.

FlashAttention reorganiza a atenção em tiles, mantém blocos em SRAM e calcula a softmax de forma incremental. O resultado é a atenção densa exata, mas com menos transferências entre HBM e memória no chip. A conexão com GEMM é direta: a mesma filosofia de blocagem e reutilização é aplicada a uma composição que inclui GEMM, máscara e softmax.

13.3 A rede feed-forward

A FFN típica contém:

com $W_1\in\mathbb{R}^{d\times d_{\text{ff}}}$ e $W_2\in\mathbb{R}^{d_{\text{ff}}\times d}$. As duas GEMMs custam:

Se $d_{\text{ff}}\approx4d$, a FFN consome aproximadamente:

Somando QKV, projeção de saída e FFN, obtemos cerca de:

FLOPs por camada, antes do termo quadrático da atenção e sem contar treinamento. Isso explica por que otimizar apenas $QK^T$ não otimiza o Transformer inteiro: projeções e FFNs são grandes consumidoras de GEMM.

Ativações como GELU, SiLU e portas, além do viés, podem ser fundidas ao epílogo da GEMM. Evitar gravar uma ativação intermediária e lê-la novamente pode valer mais do que reduzir algumas operações escalares.

14. Prefill e decode: a mesma álgebra, gargalos diferentes

Durante o prefill, muitos tokens são processados juntos. A dimensão $M=BL$ das projeções é grande. Cada bloco de pesos é reutilizado para muitas linhas de $X$, aumentando a intensidade aritmética. GEMMs grandes tendem a aproveitar bem Tensor Cores.

Durante a geração autoregressiva, o decode processa normalmente um novo token por sequência a cada passo. Para lote pequeno, a projeção se aproxima de uma multiplicação matriz-vetor:

Considere $W\in\mathbb{R}^{4096\times4096}$ em BF16. Apenas ler os pesos requer aproximadamente:

O trabalho é cerca de:

Ignorando outros dados, isso representa aproximadamente um FLOP por byte de pesos. O kernel tende a ser limitado por largura de banda. Duplicar o pico de Tensor Cores pode produzir pouco ganho se os pesos não chegarem mais rápido.

Aumentar o lote reutiliza os mesmos pesos para mais sequências e transforma GEMV em uma GEMM mais favorável. Contudo, lote maior também aumenta latência de fila e consumo de memória. Sistemas de inferência equilibram throughput e latência.

14.1 O cache KV também se move

Na atenção autoregressiva, Keys e Values anteriores ficam no KV cache. A cada token, o modelo lê partes desse cache. Conforme o contexto cresce, o volume lido por passo cresce. Multi-Query Attention e Grouped-Query Attention reduzem o número de conjuntos independentes de Keys e Values e, portanto, o tráfego e a memória do cache.

Mais uma vez, a pergunta correta não é apenas “quantas multiplicações existem?”, mas “quantas vezes cada byte atravessa a hierarquia?”.

15. Como medir GEMM sem se enganar

Um benchmark útil precisa separar alocação, inicialização e medição do kernel.

15.1 FLOP/s observado

Para uma execução com tempo $t$ em segundos:

Em GFLOP/s:

Se $\beta\neq0$, há operações adicionais sobre $C$, mas a convenção $2MNK$ facilita comparações.

15.2 Procedimento mínimo

1. aloque e inicialize fora da região medida;
2. aqueça o kernel;
3. execute várias repetições;
4. use mediana ou outra estatística robusta;
5. valide o resultado;
6. informe dimensões, tipos, layout, threads e opções de compilação;
7. meça matrizes compatíveis com a carga real;
8. em GPU, use eventos do dispositivo e sincronização correta;
9. não inclua transferências host-dispositivo sem declarar;
10. compare com uma biblioteca otimizada.

Um esqueleto de medição em CPU:

#include <algorithm>
#include <chrono>
#include <cstddef>
#include <functional>
#include <vector>

double benchmark_gemm(const std::function<void()>& kernel,
                      std::size_t M,
                      std::size_t N,
                      std::size_t K,
                      int warmups = 3,
                      int repetitions = 11) {
    for (int i = 0; i < warmups; ++i) {
        kernel();
    }

    std::vector<double> times;
    times.reserve(repetitions);

    for (int i = 0; i < repetitions; ++i) {
        const auto start = std::chrono::steady_clock::now();
        kernel();
        const auto stop = std::chrono::steady_clock::now();

        times.push_back(
            std::chrono::duration<double>(stop - start).count());
    }

    std::sort(times.begin(), times.end());
    const double median = times[times.size() / 2];
    const double flops =
        2.0 * static_cast<double>(M)
            * static_cast<double>(N)
            * static_cast<double>(K);

    return flops / median / 1.0e9;
}

Se o kernel acumula sobre $C$, restaure a mesma entrada antes de cada repetição; caso contrário, as repetições não medirão o mesmo problema. Para kernels muito pequenos, até a chamada indireta por std::function pode ser mensurável e deve ser removida do caminho cronometrado. O compilador também não pode concluir que o resultado é inútil e eliminar o trabalho: consumir ou verificar a saída resolve esse risco. Frequência dinâmica, concorrência do sistema e aquecimento térmico afetam as medições.

15.3 Correção antes de velocidade

Compare o kernel com uma referência:

O erro relativo pode ser inadequado perto de zero; por isso, uma regra comum combina tolerâncias:

Para baixa precisão, avalie também normas e impacto na aplicação. Um kernel que produz mais TFLOP/s e resultados inválidos não é uma otimização.

16. Quando usar uma biblioteca

Escrever kernels didáticos revela a estrutura da máquina. Em produção, comece com bibliotecas: uma BLAS otimizada para a CPU em uso (OpenBLAS, BLIS, MKL); cuBLAS ou cuBLASLt em CUDA, deixando a biblioteca escolher o kernel e o caminho de Tensor Cores; as bibliotecas equivalentes do acelerador escolhido, quando não for NVIDIA; e, acima de todas, compiladores e frameworks capazes de selecionar e fundir kernels automaticamente.

Essas bibliotecas embutem tudo o que este artigo construiu: microkernels escritos e ajustados por arquitetura, com os formatos da Seção 4.1; packing e blocagem multinível; paralelismo calibrado para não brigar consigo mesmo; tratamento de bordas por máscaras e kernels especializados; caminhos de precisão mista com os cuidados da Seção 12; heurísticas de seleção de kernel e, em alguns casos, autotuning; e epílogos fundidos para viés e ativação. Reproduzir esse conjunto por conta própria não é um fim de semana de trabalho; é uma carreira.

Uma implementação própria faz sentido quando a operação possui estrutura que a GEMM genérica não explora, quando precisamos fundir etapas ou quando o objetivo é aprender. Mesmo assim, a biblioteca fornece a linha de base que precisamos superar.

17. Strassen e Coppersmith–Winograd: os algoritmos rápidos que as bibliotecas não usam

Existe uma família inteira de algoritmos que multiplica matrizes com menos de $O(n^3)$ operações, e nenhum deles aparece no cuBLAS que executa os Transformers. Essa ausência não é preguiça dos engenheiros; é uma decisão técnica com razões precisas. Para entendê-la, precisamos primeiro entender o que esses algoritmos realmente fazem.

17.1 Strassen: sete multiplicações onde havia oito

Divida cada matriz em quatro blocos:

O algoritmo clássico calcula cada bloco de $C$ como a soma de dois produtos, por exemplo $C_{11}=A_{11}B_{11}+A_{12}B_{21}$: oito multiplicações de blocos e quatro adições. Em 1969, Volker Strassen mostrou que sete multiplicações bastam, ao preço de mais adições:

A leitora pode verificar qualquer uma das quatro igualdades expandindo os $M_i$; os termos indesejados cancelam-se com precisão cirúrgica. O truque não é uma multiplicação de matrizes $2\times2$: é uma receita para blocos, aplicada recursivamente. Cada um dos sete produtos é, ele próprio, uma multiplicação de matrizes pela metade do tamanho, resolvida pela mesma receita. A recorrência

resolve para $T(n)=O(n^{\log_2 7})\approx O(n^{2{,}807})$. O custo escondido está no termo $O(n^2)$: são $18$ adições ou subtrações de blocos por nível (a variante de Winograd reduz para $15$), e cada nível economiza apenas $1/8$ das multiplicações. Para reduzir as multiplicações à metade, são necessários cinco níveis de recursão, pois $(7/8)^5\approx0{,}51$, o que exige matrizes grandes, de preferência com dimensões divisíveis muitas vezes por dois. Implementações cuidadosas de Strassen chegam a vencer BLAS clássicas em CPUs, tipicamente a partir de dimensões entre algumas centenas e alguns milhares, em FP64, com espaço de trabalho extra para os blocos temporários. Guarde essas condições; elas serão o centro do argumento da Seção 17.3.

17.2 De Coppersmith–Winograd aos algoritmos galácticos

Strassen abriu uma pergunta que a matemática ainda não fechou: qual é o menor expoente $\omega$ tal que duas matrizes $n\times n$ podem ser multiplicadas em $O(n^{\omega+\epsilon})$ operações para todo $\epsilon>0$? A abordagem moderna trata a multiplicação de matrizes como um objeto algébrico, um tensor tridimensional, e busca decomposições desse tensor com poucos produtos: a receita de Strassen é exatamente uma decomposição de posto $7$ do tensor $2\times2\times2$. Coppersmith e Winograd, em 1990, aplicaram uma maquinaria muito mais pesada, o chamado método do laser sobre potências de um tensor cuidadosamente escolhido, e obtiveram $\omega<2{,}376$. Refinamentos sucessivos dessa mesma linha empurraram o limite para baixo de $2{,}372$, e conjectura-se, sem prova, que $\omega=2$.

Há um porém que o leitor de enunciados assintóticos aprende a farejar: as constantes. Os algoritmos da família Coppersmith–Winograd são chamados, sem ironia na literatura, de algoritmos galácticos: o tamanho de matriz a partir do qual venceriam o método clássico excede qualquer matriz fisicamente armazenável. Ninguém jamais executou Coppersmith–Winograd para multiplicar matrizes de verdade, e não é por falta de vontade.

A busca, ainda assim, não parou nem ficou puramente teórica. Em 2022, o AlphaTensor, da DeepMind, usou aprendizado por reforço para procurar decomposições de posto baixo e redescobriu Strassen, encontrou um algoritmo de $47$ multiplicações para blocos $4\times4$ em aritmética módulo 2, e, mais revelador para este artigo, produziu variantes otimizadas para o tempo de execução medido em uma GPU específica, não para a contagem de multiplicações. Até quem procura algoritmos novos com uma rede neural acabou admitindo que o alvo certo é o relógio, não o expoente.

17.3 Por que a prática ignora o expoente

Podemos agora justificar, uma a uma, as razões pelas quais os kernels de Transformers permanecem cúbicos. Nenhuma delas é “os engenheiros não leram os artigos”.

As dimensões são finitas e, pior, retangulares. O ganho de Strassen depende de recursão profunda, e a recursão consome a menor dimensão do problema. As GEMMs de um Transformer têm formas como $(BL)\times d\times 3d$ ou, por cabeça de atenção, $K=d_h$ de apenas $64$ ou $128$: a recursão nesse eixo termina em poucos níveis, antes que a economia de $1/8$ por nível componha algo que pague as adições extras. O regime assintótico simplesmente não é alcançado.

A baixa precisão e a estabilidade numérica caminham juntas contra as subtrações de Strassen. Expressões como $A_{12}-A_{22}$ e $B_{21}-B_{11}$ produzem cancelamentos, e a análise clássica de Higham mostra que o limite de erro de Strassen cresce aproximadamente como $n^{\log_2 12}\approx n^{3{,}58}$, contra crescimento linear em $n$ no algoritmo clássico, valendo apenas em norma, não componente a componente. Agora recorde a Seção 12: os Transformers já operam em FP16 e BF16, no limite do que o treinamento tolera, com loss scaling para resgatar gradientes do underflow. Adotar um algoritmo que amplia o erro para economizar multiplicações é resolver o problema que não temos agravando o problema que temos.

As hierarquias de memória punem exatamente a moeda em que Strassen paga. Uma adição de matrizes lê dois operandos e escreve um resultado para executar uma única operação por elemento: em FP32, um FLOP a cada doze bytes, a menor intensidade aritmética possível, irremediavelmente presa à parte inclinada do Roofline da Seção 9. Strassen troca multiplicações, que a blocagem consegue tornar limitadas por computação, por adições que nenhuma blocagem salva, e ainda exige buffers temporários para os $M_i$, pressionando as mesmas caches que a Seção 4 lutou para administrar.

O packing e o paralelismo perdem sua regularidade. A máquina de três níveis da Seção 5 depende de um passeio previsível por painéis empacotados; a árvore de recursão de Strassen fragmenta esse passeio, exige empacotar operandos derivados (somas e diferenças de blocos) em cada nível e cria um ponto de sincronização na recombinação dos sete produtos. Nada disso é impossível; tudo isso corrói margens que já eram de poucos por cento.

Os Tensor Cores mudaram a aritmética da decisão. Uma unidade matricial entrega nas GEMMs um múltiplo grande do throughput vetorial, tipicamente entre $8$ e $16$ vezes, executando precisamente a MMA clássica $D=AB+C$ sobre tiles fixos. Diante disso, os $12{,}5\%$ de multiplicações economizados por um nível de Strassen são troco, e as adições extras nem sequer executam nos Tensor Cores: ficam com as unidades vetoriais e com a memória. O hardware tornou barato exatamente aquilo que Strassen economiza, e caro exatamente aquilo que Strassen gasta.

A fusão de epílogos, da qual as Seções 1 e 13 dependem, pressupõe que cada tile de $C$ fique pronto de uma vez, para receber viés, ativação ou máscara antes de tocar a memória. Em Strassen, $C_{11}=M_1+M_4-M_5+M_7$ só está completo depois que quatro produtos independentes terminam: os intermediários precisam ser materializados, que é precisamente a leitura e escrita extra que a fusão existe para evitar.

Por fim, o decode torna a discussão inteira irrelevante para metade da inferência. A Seção 14 mostrou que a geração autoregressiva é limitada pela leitura dos pesos, com intensidade próxima de um FLOP por byte. Strassen reduz $W$, o numerador da intensidade; o gargalo do decode é $Q$, o denominador. Reduzir FLOPs de um kernel limitado por memória é afiar a faca para um duelo de pistolas.

O estudo assintótico responde “quantas operações são necessárias quando $n$ cresce sem limite?”. Este artigo responde “como manter a máquina alimentada nas matrizes que realmente executamos?”. As duas perguntas são legítimas, e a primeira produziu matemática profunda que talvez um dia mude a resposta da segunda. Para compreender Transformers como sistemas computacionais, hoje, a segunda é mais urgente.

18. Checklist mental para uma GEMM

Ao encontrar uma multiplicação matricial, pergunte:

1. Quais são $M$, $N$ e $K$?
2. As matrizes são row-major, column-major ou possuem strides?
3. É GEMM, GEMV ou uma GEMM em lote?
4. Quantos FLOPs existem?
5. Qual é o tráfego mínimo e o tráfego medido?
6. Qual é a intensidade aritmética?
7. O kernel está limitado por memória ou computação?
8. Os acessos são contíguos e coalescidos?
9. Que dados podem ser reutilizados em cache, registradores ou SRAM?
10. Os blocos cabem no nível de memória pretendido?
11. O microkernel conserva acumuladores em registradores?
12. SIMD ou Tensor Cores estão realmente sendo usados?
13. Quais são os tipos de armazenamento, produto e acumulação?
14. Há operações de epílogo que podem ser fundidas?
15. A validação numérica considera o tipo e o comprimento da redução?
16. O benchmark compara com uma biblioteca otimizada?

Se a leitora souber responder a essas perguntas, verá uma GEMM não como três laços, mas como uma coreografia entre álgebra e movimento de dados.

Conclusão

A fórmula de GEMM é simples:

Seu desempenho emerge de uma sequência de transformações:

Em CPUs, blocagem e microkernels convertem reutilização matemática em reutilização de cache e registradores. Em GPUs, tiles e memória compartilhada alimentam milhares de threads. Tensor Cores aceleram MMA, mas dependem de um pipeline capaz de entregar operandos. Precisão mista reduz bytes e aumenta throughput, ao custo de novas responsabilidades numéricas.

Nos Transformers, projeções QKV, projeção de saída e FFNs são GEMMs; a atenção contém GEMMs separadas por máscara e softmax; o prefill tende a reutilizar pesos; o decode frequentemente transmite grandes matrizes para executar relativamente pouco trabalho. A eficiência do modelo nasce menos de uma contagem abstrata de multiplicações e mais da capacidade de usar cada byte várias vezes antes que ele volte à memória lenta.

Essa é a ideia que levaremos para o restante da série: operações matematicamente equivalentes podem ter custos físicos radicalmente diferentes.

Referências

AINSLIE, J. et al. GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. EMNLP, 2023. Disponível em: https://arxiv.org/abs/2305.13245.

COPPERSMITH, D.; WINOGRAD, S. Matrix Multiplication via Arithmetic Progressions. Journal of Symbolic Computation, v. 9, n. 3, p. 251–280, 1990. DOI: https://doi.org/10.1016/S0747-7171(08)80013-2.

DAO, T.; FU, D. Y.; ERMON, S.; RUDRA, A.; RÉ, C. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. Advances in Neural Information Processing Systems, v. 35, 2022. Disponível em: https://papers.nips.cc/paper_files/paper/2022/hash/67d57c32e20fd0a7a302cb81d36e40d5-Abstract-Conference.html.

DONGARRA, J. J.; DU CROZ, J.; HAMMARLING, S.; DUFF, I. S. A Set of Level 3 Basic Linear Algebra Subprograms. ACM Transactions on Mathematical Software, v. 16, n. 1, p. 1–17, 1990.

FAWZI, A. et al. Discovering Faster Matrix Multiplication Algorithms with Reinforcement Learning. Nature, v. 610, p. 47–53, 2022. Disponível em: https://www.nature.com/articles/s41586-022-05172-4.

GOTO, K.; VAN DE GEIJN, R. A. Anatomy of High-Performance Matrix Multiplication. ACM Transactions on Mathematical Software, v. 34, n. 3, 2008. DOI: https://doi.org/10.1145/1356052.1356053.

HIGHAM, N. J. Accuracy and Stability of Numerical Algorithms. 2. ed. Philadelphia: SIAM, 2002.

INTEL. Intel Intrinsics Guide. Disponível em: https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html.

LOW, T. M.; IGUAL, F. D.; SMITH, T. M.; QUINTANA-ORTÍ, E. S. Analytical Modeling Is Enough for High-Performance BLIS. ACM Transactions on Mathematical Software, v. 43, n. 2, 2016. DOI: https://doi.org/10.1145/2925987.

MICIKEVICIUS, P. et al. Mixed Precision Training. ICLR, 2018. Disponível em: https://arxiv.org/abs/1710.03740.

NVIDIA. CUDA C++ Best Practices Guide. Seções sobre acesso coalescido, memória compartilhada e multiplicação de matrizes. Disponível em: https://docs.nvidia.com/cuda/cuda-c-best-practices-guide/.

NVIDIA. CUDA C++ Programming Guide. Seções sobre WMMA e computação de ponto flutuante. Disponível em: https://docs.nvidia.com/cuda/cuda-programming-guide/.

NVIDIA. cuBLAS Library Documentation. Seções sobre GEMM, tipos de computação e Tensor Cores. Disponível em: https://docs.nvidia.com/cuda/cublas/.

SHAZEER, N. Fast Transformer Decoding: One Write-Head Is All You Need. 2019. Disponível em: https://arxiv.org/abs/1911.02150.

STRASSEN, V. Gaussian Elimination is Not Optimal. Numerische Mathematik, v. 13, p. 354–356, 1969. DOI: https://doi.org/10.1007/BF02165411.

VASWANI, A. et al. Attention Is All You Need. Advances in Neural Information Processing Systems, v. 30, 2017. Disponível em: https://arxiv.org/abs/1706.03762.

WILLIAMS, S.; WATERMAN, A.; PATTERSON, D. Roofline: An Insightful Visual Performance Model for Multicore Architectures. Communications of the ACM, v. 52, n. 4, p. 65–76, 2009. Disponível em: https://escholarship.org/uc/item/78h8v7mr.