Phase 10 - Lesson 17

Native Sparse Attention (DeepSeek NSA)

This lesson includes a graded coding exercise that runs in your browser, unlocked with lifetime access.

Com 64k tokens, a atenção consome 70-80% da latência de decodificação. Cada laboratório de modelos abertos tem um plano para corrigir isso. O NSA da DeepSeek (melhor artigo da ACL 2025) é o que se consolidou: três ramificações de atenção paralelas — tokens comprimidos de granularidade grossa, tokens de granularidade fina retidos seletivamente e janelas deslizantes para contexto local — combinados através de um gate aprendido. Ele é alinhado ao hardware (amigável ao kernel), nativamente treinável (funciona no pré-treinamento, não acoplado na inferência) e, em decodificações de 64k, roda mais rápido que o FlashAttention ao mesmo tempo em que iguala ou supera a qualidade da atenção total. Esta lição constrói as três ramificações de ponta a ponta e mostra por que a esparsidade é diferenciável de ponta a ponta.

Tipo: Build Linguagens: Python (stdlib) Pré-requisitos: Phase 7 · 12 (KV cache, flash-attention), Phase 7 · 15 (attention variants), Phase 10 · 16 (differential attention) Tempo: ~60 minutos

Objetivos de Aprendizado

• Descrever as três ramificações de atenção da NSA e o que cada uma captura.
• Explicar por que a NSA é "nativamente treinável" enquanto os métodos anteriores de atenção esparsa eram apenas para inferência.
• Calcular a economia de computação de atenção da NSA em relação à atenção total em um contexto de 64k como uma função do tamanho do bloco de compressão e do top-k de seleção.
• Implementar a combinação de três ramificações em Python (biblioteca padrão) em uma sequência sintética curta e verificar se os pesos de gating se comportam conforme o esperado.

O Problema

A atenção total em um comprimento de sequência N custa tempo O(N^2) e O(N) de cache KV por camada. Com 64k tokens, os números de computação e largura de banda de memória são catastróficos. Estimativa teórica medida no artigo da NSA: a atenção representa de 70-80% da latência total de decodificação em 64k. Tudo o que vem a jusante (downstream) — TTFT, tokens/seg, custo por milhão de tokens — é dominado pelo custo da atenção.

A atenção esparsa é a resposta óbvia. As tentativas anteriores dividem-se em duas categorias. A esparsidade de padrão fixo (janela deslizante, espaçada/strided, bloco local) descarta informações e falha em tarefas de recuperação de longo alcance. A esparsidade em tempo de inferência (poda de cache KV, H2O, StreamingLLM) é aplicada a um modelo pré-treinado em atenção densa e recupera apenas uma fração do aceleramento potencial porque o modelo nunca foi treinado para rotear informações através do padrão esparso.

A Native Sparse Attention (Yuan et al., DeepSeek + PKU + UW, melhor artigo da ACL 2025, arXiv:2502.11089) faz as duas coisas: um padrão de esparsidade que o modelo aprende durante o pré-treinamento, implementado como um algoritmo alinhado ao kernel que realmente entrega a economia de computação na inferência. Daqui a dois anos, a NSA ou um descendente direto será a atenção padrão em todos os modelos de fronteira de contexto longo.

O Conceito

Três ramificações paralelas

Para cada consulta (query), a NSA executa a atenção três vezes, contra três visões diferentes do cache KV:

1. Ramificação comprimida (Compressed branch). Os tokens são agrupados em blocos de tamanho l (geralmente 32 ou 64). Cada bloco é comprimido em um único token de resumo por meio de uma pequena MLP aprendida. A consulta atende a esses tokens comprimidos, obtendo uma visão de granularidade grossa de toda a sequência.

2. Ramificação selecionada (Selected branch). Utilizando as pontuações de atenção da ramificação comprimida, são identificados os top-k blocos mais relevantes para a consulta atual. Os tokens de granularidade fina (não comprimidos) desses blocos são lidos e a consulta atende a todos eles. Pense na atenção da ramificação comprimida como o sinal de roteamento para a seleção.

3. Ramificação de janela deslizante (Sliding-window branch). A consulta atende aos W tokens mais recentes (geralmente 512) para contexto local. Esta ramificação captura os padrões de curto alcance baseados em estrutura (sintaxe, co-referência local) que as outras duas poderiam perder.

As saídas das três ramificações são combinadas através de um gate por posição aprendido:

out = g_cmp * out_cmp + g_sel * out_sel + g_win * out_win

g_cmp, g_sel, g_win são pesos de gate vindos de uma pequena MLP aplicada à consulta. Eles não precisam somar 1 — eles podem ponderar as ramificações de forma independente.

Por que isso é "nativamente treinável"

A etapa de seleção (blocos top-k) é discreta. Operações discretas quebram o fluxo de gradiente. Trabalhos anteriores de atenção esparsa ou pulavam a retropropagação (backpropagation) através da seleção (limitando o treinamento) ou usavam relaxações contínuas que não resultavam em esparsidade real na inferência.

A NSA contorna isso: a atenção da ramificação comprimida É uma atenção de granularidade grossa diferenciável em toda a sequência. A operação top-k apenas reutiliza as pontuações de atenção mais altas da ramificação comprimida para escolher quais blocos de granularidade fina carregar. Os gradientes fluem através das pontuações da ramificação comprimida (que influenciam tanto a saída comprimida QUANTO a lógica de seleção), e a contribuição dos blocos selecionados para a saída final também é diferenciável. A operação top_k não diferenciável é uma operação nula (no-op) no grafo computacional de avanço (forward) — ela apenas controla quais blocos são carregados da memória.

É por isso que a NSA pode ser usada no pré-treinamento de ponta a ponta. O modelo aprende a rotear informações pelas três ramificações em conjunto, produzindo um padrão esparso que, na inferência, realmente entrega a aceleração prometida.

Kernel alinhado ao hardware

O kernel da NSA é projetado para hierarquias de memória de GPUs modernas. O kernel carrega consultas por grupos de GQA (loop externo), busca os blocos KV esparsos correspondentes por grupo (loop interno) e executa a atenção na SRAM. Como cada grupo de consultas vê os mesmos blocos selecionados (a seleção é por grupo de consultas, não por cabeça de consulta), os carregamentos de KV são amortizados ao longo do grupo. A intensidade aritmética permanece alta.

O artigo relata kernels Triton rodando 9x mais rápido que o FlashAttention em decodificações de 64k, com a taxa de aceleração aumentando com o comprimento da sequência. Kernels para passos de avanço (forward) e retropropagação (backward) são fornecidos.

O orçamento de computação

Seja N o comprimento da sequência, l o tamanho do bloco de compressão, k a contagem de seleção top-k, w a janela deslizante e b o tamanho do bloco selecionado (normalmente igual a l).

• Ramificação comprimida: O(N/l) chaves por consulta, totalizando O(N * N / l).
• Ramificação selecionada: O(k * b) chaves por consulta, totalizando O(N * k * b).
• Ramificação de janela deslizante: O(w) chaves por consulta, totalizando O(N * w).

Total: O(N * (N/l + k*b + w)).

Com N = 64k, l = 64, k = 16, b = 64, w = 512: o custo por consulta é 1000 + 1024 + 512 = 2536 chaves. A atenção total é de 64000 chaves. Uma redução computacional de 25x.

Com N = 128k, l = 64, k = 16, b = 64, w = 512: o custo por consulta é 2000 + 1024 + 512 = 3536 chaves. A atenção total é de 128000 chaves. Redução de 36x. O benefício cresce com o comprimento da sequência, o que é o objetivo principal.

Como ela se compara

O MoBA (Moonshot, arXiv:2502.13189) foi publicado simultaneamente e adota uma abordagem semelhante de "três é melhor do que um", aplicando o princípio de MoE a blocos de atenção. NSA e MoBA são as duas arquiteturas fundamentais para se conhecer no pré-treinamento de contexto longo em 2026.

Construa

O arquivo code/main.py implementa as três ramificações em uma sequência sintética curta e mostra:

• A MLP de compressão (uma linha de base simples de pooling médio (mean-pool) é usada para clareza pedagógica; a NSA real usa uma MLP aprendida).
• A seleção de blocos top-k impulsionada pelas pontuações da ramificação comprimida.
• A atenção de janela deslizante nos últimos w tokens.
• A combinação por gates (gated).
• Uma impressão de contagem de computação comparando com a atenção total.

Passo 1: comprimir tokens em blocos

def compress(K, l):
    n = len(K)
    n_blocks = (n + l - 1) // l
    out = []
    for b in range(n_blocks):
        start, end = b * l, min((b + 1) * l, n)
        block = K[start:end]
        summary = [sum(row[d] for row in block) / len(block) for d in range(len(K[0]))]
        out.append(summary)
    return out

Passo 2: atenção da ramificação comprimida

Execute a atenção softmax da consulta em relação às chaves comprimidas. As pontuações da ramificação comprimida servem também como o sinal para a seleção top-k.

Passo 3: seleção de blocos top-k

Escolha os índices dos k blocos comprimidos de maior pontuação. Carregue os tokens originais não comprimidos desses blocos e execute a atenção neles.

Passo 4: atenção de janela deslizante

Pegue os últimos w tokens e execute a atenção padrão contra eles.

Passo 5: gate + combinar

Uma pequena MLP aplicada à consulta produz três pesos de gate. A saída final é a soma ponderada das saídas das três ramificações.

Passo 6: contagem de computação

Imprima o número de chaves atendidas por consulta para cada ramificação e o total. Compare com N (atenção total). Em um exemplo sintético de 1024 tokens com l = 32, k = 4, w = 128, a NSA vê 32 + 128 + 128 = 288 chaves por consulta em comparação com 1024 para atenção total — 3.5x menos.

Use

A NSA está sendo enviada no próprio pipeline de pré-treinamento de contexto longo da DeepSeek. Status de integração em pilhas de inferência públicas a partir de abril de 2026:

• DeepSeek interno: nativo, os pesos publicados usam NSA ou seu sucessor DSA (Deepseek Sparse Attention).
• vLLM: suporte experimental a NSA em desenvolvimento para pesos do DeepSeek-V3.x.
• SGLang: benchmarks de NSA publicados; caminho de produção segue o vLLM.
• llama.cpp / CPU: não suportado; o overhead da decomposição do kernel não compensa com o rendimento (throughput) de CPU.

Quando escolher a NSA:

• Execução de pré-treinamento ou treinamento contínuo visando um contexto de mais de 64k com um orçamento de computação sério.
• Inferência dos próprios checkpoints de contexto longo da DeepSeek. Os pesos são nativos do NSA.

Quando não escolher:

• Servir um modelo pré-treinado de atenção densa existente. Não é possível adaptar a NSA sem treinamento contínuo.
• Contexto abaixo de 16k. O overhead das três ramificações domina a economia.
• Chat interativo com lote 1 (Batch-1). A decodificação sensível à latência se beneficia, mas apenas em contextos longos.

Envie

Esta lição produz outputs/skill-nsa-integrator.md. Dada uma especificação de execução de pré-treinamento de contexto longo, ela produz um plano de integração da NSA: tamanho do bloco de compressão, top-k, janela deslizante, largura da MLP do gate, escolha do kernel e as avaliações específicas de contexto longo que justificariam a mudança de arquitetura.

Exercícios

1. Execute code/main.py em uma sequência sintética de 1024 tokens. Faça uma varredura (sweep) de (l, k, w) em três predefinições e imprima as contagens de computação. Identifique a predefinição que alcança a menor contagem de chaves por consulta enquanto mantém 95% de recuperação (recall) em relação à atenção total em um teste do tipo "agulha no palheiro" (needle-in-a-haystack).

2. Substitua o compressor de pooling médio (mean-pool) por uma pequena MLP aprendida (2 camadas, 32 unidades ocultas). Treine-a em uma tarefa sintética onde o sinal é a média de um bloco. Meça a diferença (gap) de perplexidade em relação à linha de base do mean-pool em dados de validação.

3. Implemente a MLP de gate. Ela recebe a consulta como entrada e produz três escalares. Mostre que o gate se comporta de maneira sensata: ponderação quase uniforme em consultas aleatórias, peso alto na ramificação selecionada quando a consulta atinge um bloco muito anterior.

4. Calcule o orçamento de memória de cache KV para um modelo de 70B com suporte a NSA em um contexto de 128k. As cabeças de KV são 8, a dimensão da cabeça é 128, BF16. Compare com a atenção total e com a MLA (a Fase 10 · 14 mostrou os números da MLA). Identifique o comprimento da sequência onde o cache KV da ramificação de granularidade fina da NSA se iguala à atenção total.

5. Leia a Seção 4 do artigo da NSA (arXiv:2502.11089) e explique em três frases por que as pontuações de atenção da ramificação comprimida são reutilizadas para a seleção top-k em vez de computar uma pontuação de roteamento separada. Associe a resposta ao fluxo de gradientes.

Termos-Chave

Leituras Adicionais

• Yuan et al. — Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention (arXiv:2502.11089, ACL 2025 Best Paper) — o artigo
• DeepSeek-V3 Technical Report (arXiv:2412.19437) — a família de arquitetura que a NSA visa
• Moonshot AI — MoBA: Mixture of Block Attention for Long-Context LLMs (arXiv:2502.13189) — trabalho concorrente, atenção no estilo MoE sobre blocos
• Beltagy et al. — Longformer: The Long-Document Transformer (arXiv:2004.05150) — origens da janela deslizante
• Xiao et al. — StreamingLLM: Efficient Streaming Language Models with Attention Sinks (arXiv:2309.17453) — linha de base de esparsidade em tempo de inferência que a NSA aprimora
• Dao et al. — FlashAttention-2 (arXiv:2307.08691) — a linha de base de atenção total que os kernels da NSA superam em 64k